Upload folder using huggingface_hub

seq_10_150_3_3_encode_occ_separate_offsets_bs8_128_no_seqindex_long_lastvalid_ablation_grid/backups/configs/experiments/ablate_grid_tokens.yaml

@@ -0,0 +1,106 @@
+# Config format schema number, the yaml support to valid case source from different dataset
+time_info: &time_info
+  num_historical_steps: 11
+  num_future_steps: 80
+  use_intention: True
+  token_size: 2048
+  predict_motion: True
+  predict_state: True
+  predict_map: True
+  predict_occ: True
+  state_token:
+    invalid: 0
+    valid: 1
+    enter: 2
+    exit: 3
+  pl2seed_radius: 75.
+  disable_grid_token: True
+  grid_range: 150. # 2 times of pl2seed_radius
+  grid_interval: 3.
+  angle_interval: 3.
+  seed_size: 1
+  buffer_size: 128
+  max_num: 32
+
+Dataset:
+  root:
+  train_batch_size: 1
+  val_batch_size: 1
+  test_batch_size: 1
+  shuffle: True
+  num_workers: 1
+  pin_memory: True
+  persistent_workers: True
+  train_raw_dir: 'data/waymo_processed/training'
+  val_raw_dir: 'data/waymo_processed/validation'
+  test_raw_dir: 'data/waymo_processed/validation'
+  val_tfrecords_splitted: data/waymo_processed/training/validation_tfrecords_splitted
+  transform: WaymoTargetBuilder
+  train_processed_dir:
+  val_processed_dir:
+  test_processed_dir:
+  dataset: 'scalable'
+  <<: *time_info
+
+Trainer:
+  strategy: ddp_find_unused_parameters_false
+  accelerator: 'gpu'
+  devices: 1
+  max_epochs: 32
+  save_ckpt_path:
+  num_nodes: 1
+  mode:
+  ckpt_path:
+  precision: 32
+  accumulate_grad_batches: 1
+  overfit_epochs: 6000
+
+Model:
+  predictor: 'smart'
+  decoder_type: 'agent_decoder' # choose from ['agent_decoder', 'occ_decoder']
+  dataset: 'waymo'
+  input_dim: 2
+  hidden_dim: 128
+  output_dim: 2
+  output_head: False
+  num_heads: 8
+  <<: *time_info
+  head_dim: 16
+  dropout: 0.1
+  num_freq_bands: 64
+  lr: 0.0005
+  warmup_steps: 0
+  total_steps: 32
+  predict_map_token: False
+  num_recurrent_steps_val: 300
+  val_open_loop: False
+  val_close_loop: True
+  val_insert: False
+  n_rollout_close_val: 1
+  decoder:
+    <<: *time_info
+    num_map_layers: 3
+    num_agent_layers: 6
+    a2a_radius: 60
+    pl2pl_radius: 10
+    pl2a_radius: 30
+    a2sa_radius: 10
+    pl2sa_radius: 10
+    time_span: 60
+  loss_weight:
+    token_cls_loss: 1
+    map_token_loss: 1
+    state_cls_loss: 10
+    type_cls_loss: 5
+    pos_cls_loss: 1
+    head_cls_loss: 1
+    offset_reg_loss: 5
+    shape_reg_loss: .2
+    pos_reg_loss: 10
+    state_weight: [0.1, 0.1, 0.8] # invalid, valid, exit
+    seed_state_weight: [0.1, 0.9] # invalid, enter
+    seed_type_weight: [0.8, 0.1, 0.1]
+    agent_occ_pos_weight: 100
+    pt_occ_pos_weight: 5
+    agent_occ_loss: 10
+    pt_occ_loss: 10

seq_10_150_3_3_encode_occ_separate_offsets_bs8_128_no_seqindex_long_lastvalid_ablation_grid/backups/configs/ours_long_term.yaml

@@ -0,0 +1,105 @@
+# Config format schema number, the yaml support to valid case source from different dataset
+time_info: &time_info
+  num_historical_steps: 11
+  num_future_steps: 80
+  use_intention: True
+  token_size: 2048
+  predict_motion: True
+  predict_state: True
+  predict_map: True
+  predict_occ: True
+  state_token:
+    invalid: 0
+    valid: 1
+    enter: 2
+    exit: 3
+  pl2seed_radius: 75.
+  grid_range: 150. # 2 times of pl2seed_radius
+  grid_interval: 3.
+  angle_interval: 3.
+  seed_size: 1
+  buffer_size: 128
+  max_num: 32
+
+Dataset:
+  root:
+  train_batch_size: 1
+  val_batch_size: 1
+  test_batch_size: 1
+  shuffle: True
+  num_workers: 1
+  pin_memory: True
+  persistent_workers: True
+  train_raw_dir: 'data/waymo_processed/training'
+  val_raw_dir: 'data/waymo_processed/validation'
+  test_raw_dir: 'data/waymo_processed/validation'
+  val_tfrecords_splitted: data/waymo_processed/training/validation_tfrecords_splitted
+  transform: WaymoTargetBuilder
+  train_processed_dir:
+  val_processed_dir:
+  test_processed_dir:
+  dataset: 'scalable'
+  <<: *time_info
+
+Trainer:
+  strategy: ddp_find_unused_parameters_false
+  accelerator: 'gpu'
+  devices: 1
+  max_epochs: 32
+  save_ckpt_path:
+  num_nodes: 1
+  mode:
+  ckpt_path:
+  precision: 32
+  accumulate_grad_batches: 1
+  overfit_epochs: 6000
+
+Model:
+  predictor: 'smart'
+  decoder_type: 'agent_decoder' # choose from ['agent_decoder', 'occ_decoder']
+  dataset: 'waymo'
+  input_dim: 2
+  hidden_dim: 128
+  output_dim: 2
+  output_head: False
+  num_heads: 8
+  <<: *time_info
+  head_dim: 16
+  dropout: 0.1
+  num_freq_bands: 64
+  lr: 0.0005
+  warmup_steps: 0
+  total_steps: 32
+  predict_map_token: False
+  num_recurrent_steps_val: 300
+  val_open_loop: False
+  val_close_loop: True
+  val_insert: False
+  n_rollout_close_val: 1
+  decoder:
+    <<: *time_info
+    num_map_layers: 3
+    num_agent_layers: 6
+    a2a_radius: 60
+    pl2pl_radius: 10
+    pl2a_radius: 30
+    a2sa_radius: 10
+    pl2sa_radius: 10
+    time_span: 60
+  loss_weight:
+    token_cls_loss: 1
+    map_token_loss: 1
+    state_cls_loss: 10
+    type_cls_loss: 5
+    pos_cls_loss: 1
+    head_cls_loss: 1
+    offset_reg_loss: 5
+    shape_reg_loss: .2
+    pos_reg_loss: 10
+    state_weight: [0.1, 0.1, 0.8] # invalid, valid, exit
+    seed_state_weight: [0.1, 0.9] # invalid, enter
+    seed_type_weight: [0.8, 0.1, 0.1]
+    agent_occ_pos_weight: 100
+    pt_occ_pos_weight: 5
+    agent_occ_loss: 10
+    pt_occ_loss: 10

seq_10_150_3_3_encode_occ_separate_offsets_bs8_128_no_seqindex_long_lastvalid_ablation_grid/backups/configs/ours_standard.yaml

@@ -0,0 +1,101 @@
+# Config format schema number, the yaml support to valid case source from different dataset
+time_info: &time_info
+  num_historical_steps: 11
+  num_future_steps: 80
+  use_intention: True
+  token_size: 2048
+  predict_motion: True
+  predict_state: True
+  predict_map: False
+  predict_occ: True
+  state_token:
+    invalid: 0
+    valid: 1
+    enter: 2
+    exit: 3
+  pl2seed_radius: 75.
+  grid_range: 150. # 2 times of pl2seed_radius
+  grid_interval: 3.
+  angle_interval: 3.
+  seed_size: 1
+  buffer_size: 128
+
+Dataset:
+  root:
+  train_batch_size: 1
+  val_batch_size: 1
+  test_batch_size: 1
+  shuffle: True
+  num_workers: 1
+  pin_memory: True
+  persistent_workers: True
+  train_raw_dir: ["data/waymo_processed/training"]
+  val_raw_dir: ["data/waymo_processed/validation"]
+  test_raw_dir: ["data/waymo_processed/validation"]
+  transform: WaymoTargetBuilder
+  train_processed_dir:
+  val_processed_dir:
+  test_processed_dir:
+  dataset: "scalable"
+  <<: *time_info
+
+Trainer:
+  strategy: ddp_find_unused_parameters_false
+  accelerator: "gpu"
+  devices: 1
+  max_epochs: 32
+  overfit_epochs: 6000
+  save_ckpt_path:
+  num_nodes: 1
+  mode:
+  ckpt_path:
+  precision: 32
+  accumulate_grad_batches: 1
+
+Model:
+  predictor: "smart"
+  decoder_type: "agent_decoder" # choose from ['agent_decoder', 'occ_decoder']
+  dataset: "waymo"
+  input_dim: 2
+  hidden_dim: 128
+  output_dim: 2
+  output_head: False
+  num_heads: 8
+  <<: *time_info
+  head_dim: 16
+  dropout: 0.1
+  num_freq_bands: 64
+  lr: 0.0005
+  warmup_steps: 0
+  total_steps: 32
+  predict_map_token: False
+  num_recurrent_steps_val: -1
+  val_open_loop: True
+  val_close_loop: False
+  val_insert: False
+  decoder:
+    <<: *time_info
+    num_map_layers: 3
+    num_agent_layers: 6
+    a2a_radius: 60
+    pl2pl_radius: 10
+    pl2a_radius: 30
+    a2sa_radius: 10
+    pl2sa_radius: 10
+    time_span: 60
+  loss_weight:
+    token_cls_loss: 1
+    map_token_loss: 1
+    state_cls_loss: 10
+    type_cls_loss: 5
+    pos_cls_loss: 1
+    head_cls_loss: 1
+    offset_reg_loss: 5
+    shape_reg_loss: .2
+    state_weight: [0.1, 0.1, 0.8] # invalid, valid, exit
+    seed_state_weight: [0.1, 0.9] # invalid, enter
+    seed_type_weight: [0.8, 0.1, 0.1]
+    agent_occ_pos_weight: 100
+    pt_occ_pos_weight: 5
+    agent_occ_loss: 10
+    pt_occ_loss: 10

seq_10_150_3_3_encode_occ_separate_offsets_bs8_128_no_seqindex_long_lastvalid_ablation_grid/backups/configs/ours_standard_decode_occ.yaml

@@ -0,0 +1,100 @@
+# Config format schema number, the yaml support to valid case source from different dataset
+time_info: &time_info
+  num_historical_steps: 11
+  num_future_steps: 80
+  use_intention: True
+  token_size: 2048
+  predict_motion: False
+  predict_state: False
+  predict_map: False
+  predict_occ: True
+  state_token:
+    invalid: 0
+    valid: 1
+    enter: 2
+    exit: 3
+  pl2seed_radius: 75.
+  grid_range: 150. # 2 times of pl2seed_radius
+  grid_interval: 3.
+  angle_interval: 3.
+  seed_size: 1
+  buffer_size: 128
+
+Dataset:
+  root:
+  train_batch_size: 1
+  val_batch_size: 1
+  test_batch_size: 1
+  shuffle: True
+  num_workers: 1
+  pin_memory: True
+  persistent_workers: True
+  train_raw_dir: ["data/waymo_processed/training"]
+  val_raw_dir: ["data/waymo_processed/validation"]
+  test_raw_dir: ["data/waymo_processed/validation"]
+  transform: WaymoTargetBuilder
+  train_processed_dir:
+  val_processed_dir:
+  test_processed_dir:
+  dataset: "scalable"
+  <<: *time_info
+
+Trainer:
+  strategy: ddp_find_unused_parameters_false
+  accelerator: "gpu"
+  devices: 1
+  max_epochs: 32
+  overfit_epochs: 6000
+  save_ckpt_path:
+  num_nodes: 1
+  mode:
+  ckpt_path:
+  precision: 32
+  accumulate_grad_batches: 1
+
+Model:
+  predictor: "smart"
+  decoder_type: "occ_decoder" # choose from ['agent_decoder', 'occ_decoder']
+  dataset: "waymo"
+  input_dim: 2
+  hidden_dim: 128
+  output_dim: 2
+  output_head: False
+  num_heads: 8
+  <<: *time_info
+  head_dim: 16
+  dropout: 0.1
+  num_freq_bands: 64
+  lr: 0.0005
+  warmup_steps: 0
+  total_steps: 32
+  predict_map_token: False
+  num_recurrent_steps_val: -1
+  val_open_loop: True
+  val_closed_loop: False
+  decoder:
+    <<: *time_info
+    num_map_layers: 3
+    num_agent_layers: 6
+    a2a_radius: 60
+    pl2pl_radius: 10
+    pl2a_radius: 30
+    a2sa_radius: 10
+    pl2sa_radius: 10
+    time_span: 60
+  loss_weight:
+    token_cls_loss: 1
+    map_token_loss: 1
+    state_cls_loss: 10
+    type_cls_loss: 5
+    pos_cls_loss: 1
+    head_cls_loss: 1
+    offset_reg_loss: 5
+    shape_reg_loss: .2
+    state_weight: [0.1, 0.1, 0.8] # invalid, valid, exit
+    seed_state_weight: [0.1, 0.9] # invalid, enter
+    seed_type_weight: [0.8, 0.1, 0.1]
+    agent_occ_pos_weight: 100
+    pt_occ_pos_weight: 5
+    agent_occ_loss: 10
+    pt_occ_loss: 10

seq_10_150_3_3_encode_occ_separate_offsets_bs8_128_no_seqindex_long_lastvalid_ablation_grid/backups/configs/pretrain_scalable_map.yaml

@@ -0,0 +1,97 @@
+# Config format schema number, the yaml support to valid case source from different dataset
+time_info: &time_info
+  num_historical_steps: 11
+  num_future_steps: 80
+  use_intention: True
+  token_size: 2048
+  predict_motion: False
+  predict_state: False
+  predict_map: True
+  predict_occ: False
+  state_token:
+    invalid: 0
+    valid: 1
+    enter: 2
+    exit: 3
+  pl2seed_radius: 75.
+  grid_range: 150. # 2 times of pl2seed_radius
+  grid_interval: 3.
+  angle_interval: 3.
+  seed_size: 1
+  buffer_size: 32
+
+Dataset:
+  root:
+  train_batch_size: 1
+  val_batch_size: 1
+  test_batch_size: 1
+  shuffle: True
+  num_workers: 1
+  pin_memory: True
+  persistent_workers: True
+  train_raw_dir: ["data/waymo_processed/training"]
+  val_raw_dir: ["data/waymo_processed/validation"]
+  test_raw_dir: ["data/waymo_processed/validation"]
+  transform: WaymoTargetBuilder
+  train_processed_dir:
+  val_processed_dir:
+  test_processed_dir:
+  dataset: "scalable"
+  <<: *time_info
+
+Trainer:
+  strategy: ddp_find_unused_parameters_false
+  accelerator: "gpu"
+  devices: 1
+  max_epochs: 32
+  overfit_epochs: 6000
+  save_ckpt_path:
+  num_nodes: 1
+  mode:
+  ckpt_path:
+  precision: 32
+  accumulate_grad_batches: 1
+
+Model:
+  mode: "train"
+  predictor: "smart"
+  decoder_type: "agent_decoder"
+  dataset: "waymo"
+  input_dim: 2
+  hidden_dim: 128
+  output_dim: 2
+  output_head: False
+  num_heads: 8
+  <<: *time_info
+  head_dim: 16
+  dropout: 0.1
+  num_freq_bands: 64
+  lr: 0.0005
+  warmup_steps: 0
+  total_steps: 32
+  predict_map_token: False
+  decoder:
+    <<: *time_info
+    num_map_layers: 3
+    num_agent_layers: 6
+    a2a_radius: 60
+    pl2pl_radius: 10
+    pl2a_radius: 30
+    a2sa_radius: 10
+    pl2sa_radius: 10
+    time_span: 60
+  loss_weight:
+    token_cls_loss: 1
+    map_token_loss: 1
+    state_cls_loss: 10
+    type_cls_loss: 5
+    pos_cls_loss: 1
+    head_cls_loss: 1
+    shape_reg_loss: .2
+    state_weight: [0.1, 0.1, 0.8] # invalid, valid, exit
+    seed_state_weight: [0.1, 0.9] # invalid, enter
+    seed_type_weight: [0.8, 0.1, 0.1]
+    agent_occ_pos_weight: 100
+    pt_occ_pos_weight: 5
+    agent_occ_loss: 10
+    pt_occ_loss: 10

seq_10_150_3_3_encode_occ_separate_offsets_bs8_128_no_seqindex_long_lastvalid_ablation_grid/backups/configs/smart.yaml

@@ -0,0 +1,70 @@
+# Config format schema number, the yaml support to valid case source from different dataset
+time_info: &time_info
+  num_historical_steps: 11
+  num_future_steps: 80
+  use_intention: True
+  token_size: 2048
+  disable_invalid: True
+  use_special_motion_token: False
+  use_state_token: False
+  only_state: False
+
+Dataset:
+  root:
+  train_batch_size: 1
+  val_batch_size: 1
+  test_batch_size: 1
+  shuffle: True
+  num_workers: 1
+  pin_memory: True
+  persistent_workers: True
+  train_raw_dir: ["data/waymo_processed/training"]
+  val_raw_dir: ["data/waymo_processed/validation"]
+  test_raw_dir: ["data/waymo_processed/validation"]
+  transform: WaymoTargetBuilder
+  train_processed_dir:
+  val_processed_dir:
+  test_processed_dir:
+  dataset: "scalable"
+  <<: *time_info
+
+Trainer:
+  strategy: ddp_find_unused_parameters_false
+  accelerator: "gpu"
+  devices: 1
+  max_epochs: 32
+  overfit_epochs: 5000
+  save_ckpt_path:
+  num_nodes: 1
+  mode:
+  ckpt_path:
+  precision: 32
+  accumulate_grad_batches: 1
+
+Model:
+  mode: "train"
+  predictor: "smart"
+  dataset: "waymo"
+  input_dim: 2
+  hidden_dim: 128
+  output_dim: 2
+  output_head: False
+  num_heads: 8
+  <<: *time_info
+  head_dim: 16
+  dropout: 0.1
+  num_freq_bands: 64
+  lr: 0.0005
+  warmup_steps: 0
+  total_steps: 32
+  decoder:
+    <<: *time_info
+    num_map_layers: 3
+    num_agent_layers: 6
+    a2a_radius: 60
+    pl2pl_radius: 10
+    pl2a_radius: 30
+    time_span: 30
+  loss_weight:
+    token_cls_loss: 1
+    state_cls_loss: 5

seq_10_150_3_3_encode_occ_separate_offsets_bs8_128_no_seqindex_long_lastvalid_ablation_grid/backups/data_preprocess.py

@@ -0,0 +1,916 @@
+# Not a contribution
+# Changes made by NVIDIA CORPORATION & AFFILIATES enabling <CAT-K> or otherwise documented as
+# NVIDIA-proprietary are not a contribution and subject to the following terms and conditions:
+# SPDX-FileCopyrightText: Copyright (c) <year> NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: LicenseRef-NvidiaProprietary
+#
+# NVIDIA CORPORATION, its affiliates and licensors retain all intellectual
+# property and proprietary rights in and to this material, related
+# documentation and any modifications thereto. Any use, reproduction,
+# disclosure or distribution of this material and related documentation
+# without an express license agreement from NVIDIA CORPORATION or
+# its affiliates is strictly prohibited.
+
+import signal
+import multiprocessing
+import os
+import numpy as np
+import pandas as pd
+import tensorflow as tf
+import torch
+import pickle
+import easydict
+from functools import partial
+from scipy.interpolate import interp1d
+from argparse import ArgumentParser
+from tqdm import tqdm
+from typing import Any, Dict, List, Optional
+from waymo_open_dataset.protos import scenario_pb2
+
+
+MIN_VALID_STEPS = 15
+
+
+_polygon_types = ['VEHICLE', 'BIKE', 'BUS', 'PEDESTRIAN']
+_polygon_light_type = ['LANE_STATE_STOP', 'LANE_STATE_GO', 'LANE_STATE_CAUTION', 'LANE_STATE_UNKNOWN']
+_point_types = ['DASH_SOLID_YELLOW', 'DASH_SOLID_WHITE', 'DASHED_WHITE', 'DASHED_YELLOW',
+                'DOUBLE_SOLID_YELLOW', 'DOUBLE_SOLID_WHITE', 'DOUBLE_DASH_YELLOW', 'DOUBLE_DASH_WHITE',
+                'SOLID_YELLOW', 'SOLID_WHITE', 'SOLID_DASH_WHITE', 'SOLID_DASH_YELLOW', 'EDGE',
+                'NONE', 'UNKNOWN', 'CROSSWALK', 'CENTERLINE']
+_polygon_to_polygon_types = ['NONE', 'PRED', 'SUCC', 'LEFT', 'RIGHT']
+
+
+Lane_type_hash = {
+    4: "BIKE",
+    3: "VEHICLE",
+    2: "VEHICLE",
+    1: "BUS"
+}
+
+boundary_type_hash = {
+        5: "UNKNOWN",
+        6: "DASHED_WHITE",
+        7: "SOLID_WHITE",
+        8: "DOUBLE_DASH_WHITE",
+        9: "DASHED_YELLOW",
+        10: "DOUBLE_DASH_YELLOW",
+        11: "SOLID_YELLOW",
+        12: "DOUBLE_SOLID_YELLOW",
+        13: "DASH_SOLID_YELLOW",
+        14: "UNKNOWN",
+        15: "EDGE",
+        16: "EDGE"
+}
+
+
+def safe_list_index(ls: List[Any], elem: Any) -> Optional[int]:
+    try:
+        return ls.index(elem)
+    except ValueError:
+        return None
+
+
+# def get_agent_features(df: pd.DataFrame, av_id, num_historical_steps=11, dim=3, num_steps=91) -> Dict[str, Any]:
+#     if args.disable_invalid:  # filter out agents that are unseen during the historical time steps
+#         historical_df = df[df['timestep'] == num_historical_steps-1] # extract the timestep==10 (current)
+#         agent_ids = list(historical_df['track_id'].unique()) # these agents are seen at timestep==10 (current)
+#         df = df[df['track_id'].isin(agent_ids)] # remove other agents
+#     else:
+#         agent_ids = list(df['track_id'].unique())
+
+#     num_agents = len(agent_ids)
+#     # initialization
+#     valid_mask              = torch.zeros(num_agents, num_steps, dtype=torch.bool)
+#     current_valid_mask      = torch.zeros(num_agents, dtype=torch.bool)
+#     predict_mask            = torch.zeros(num_agents, num_steps, dtype=torch.bool)
+#     agent_id: List[Optional[str]] = [None] * num_agents
+#     agent_type              = torch.zeros(num_agents, dtype=torch.uint8)
+#     agent_category          = torch.zeros(num_agents, dtype=torch.uint8)
+#     position                = torch.zeros(num_agents, num_steps, dim, dtype=torch.float)
+#     heading                 = torch.zeros(num_agents, num_steps, dtype=torch.float)
+#     velocity                = torch.zeros(num_agents, num_steps, dim, dtype=torch.float)
+#     shape                   = torch.zeros(num_agents, num_steps, dim, dtype=torch.float)
+
+#     for track_id, track_df in df.groupby('track_id'):
+#         agent_idx = agent_ids.index(track_id)
+#         all_agent_steps = track_df['timestep'].values
+#         valid_agent_steps = all_agent_steps[track_df['validity'].astype(np.bool_)].astype(np.int32)
+#         valid_mask[agent_idx, valid_agent_steps] = True
+#         current_valid_mask[agent_idx] = valid_mask[agent_idx, num_historical_steps - 1] # current timestep 10
+#         if args.disable_invalid:
+#             predict_mask[agent_idx, valid_agent_steps] = True
+#         else:
+#             predict_mask[agent_idx] = True
+#         predict_mask[agent_idx, :num_historical_steps] = False
+#         if not current_valid_mask[agent_idx]:
+#             predict_mask[agent_idx, num_historical_steps:] = False
+
+#         # TODO: why using vector_repr?
+#         if vector_repr:  # a time step t is valid only when both t and t-1 are valid
+#             valid_mask[agent_idx, 1 : num_historical_steps] = (
+#                 valid_mask[agent_idx, : num_historical_steps - 1] &
+#                 valid_mask[agent_idx, 1 : num_historical_steps])
+#             valid_mask[agent_idx, 0] = False
+
+#         agent_id[agent_idx] = track_id
+#         agent_type[agent_idx] = _agent_types.index(track_df['object_type'].values[0])
+#         agent_category[agent_idx] = track_df['object_category'].values[0]
+#         position[agent_idx, valid_agent_steps, :3] = torch.from_numpy(np.stack([track_df['position_x'].values[valid_agent_steps],
+#                                                                                 track_df['position_y'].values[valid_agent_steps],
+#                                                                                 track_df['position_z'].values[valid_agent_steps]],
+#                                                                         axis=-1)).float()
+#         heading[agent_idx, valid_agent_steps] = torch.from_numpy(track_df['heading'].values[valid_agent_steps]).float()
+#         velocity[agent_idx, valid_agent_steps, :2] = torch.from_numpy(np.stack([track_df['velocity_x'].values[valid_agent_steps],
+#                                                                                 track_df['velocity_y'].values[valid_agent_steps]],
+#                                                                         axis=-1)).float()
+#         shape[agent_idx, valid_agent_steps, :3] = torch.from_numpy(np.stack([track_df['length'].values[valid_agent_steps],
+#                                                                              track_df['width'].values[valid_agent_steps],
+#                                                                              track_df["height"].values[valid_agent_steps]],
+#                                                                         axis=-1)).float()
+#     av_idx = agent_id.index(av_id)
+#     if split == 'test':
+#         predict_mask[current_valid_mask
+#                      | (agent_category == 2)
+#                      | (agent_category == 3), num_historical_steps:] = True
+
+#     return {
+#         'num_nodes': num_agents,
+#         'av_index': av_idx,
+#         'valid_mask': valid_mask,
+#         'predict_mask': predict_mask,
+#         'id': agent_id,
+#         'type': agent_type,
+#         'category': agent_category,
+#         'position': position,
+#         'heading': heading,
+#         'velocity': velocity,
+#         'shape': shape
+#     }
+
+
+def get_agent_features(track_infos: Dict[str, np.ndarray], av_id: int, num_historical_steps: int, num_steps: int) -> Dict[str, Any]:
+
+    agent_idx_to_add = []
+    for i in range(len(track_infos['object_id'])):
+        is_visible = track_infos['valid'][i, num_historical_steps - 1]
+        valid_steps = np.where(track_infos['valid'][i])[0]
+        valid_start, valid_end = valid_steps[0], valid_steps[-1]
+        is_valid = (valid_end - valid_start + 1) >= MIN_VALID_STEPS
+
+        if (is_visible or not args.disable_invalid) and is_valid:
+            agent_idx_to_add.append(i)
+
+    num_agents = len(agent_idx_to_add)
+    out_dict = {
+        'num_nodes':  num_agents,
+        'valid_mask': torch.zeros(num_agents, num_steps, dtype=torch.bool),
+        'role':       torch.zeros(num_agents, 3, dtype=torch.bool),
+        'id':         torch.zeros(num_agents, dtype=torch.int64) - 1,
+        'type':       torch.zeros(num_agents, dtype=torch.uint8),
+        'category':   torch.zeros(num_agents, dtype=torch.uint8),
+        'position':   torch.zeros(num_agents, num_steps, 3, dtype=torch.float),
+        'heading':    torch.zeros(num_agents, num_steps, dtype=torch.float),
+        'velocity':   torch.zeros(num_agents, num_steps, 2, dtype=torch.float),
+        'shape':      torch.zeros(num_agents, num_steps, 3, dtype=torch.float),
+    }
+
+    for i, idx in enumerate(agent_idx_to_add):
+
+        out_dict['role'][i] = torch.from_numpy(track_infos['role'][idx])
+        out_dict['id'][i] = track_infos['object_id'][idx]
+        out_dict['type'][i] = track_infos['object_type'][idx]
+        out_dict['category'][i] = idx in track_infos['tracks_to_predict']
+
+        valid = track_infos["valid"][idx]  # [n_step]
+        states = track_infos["states"][idx]
+
+        object_shape = states[:, 3:6]  # [n_step, 3], length, width, height
+        object_shape = object_shape[valid].mean(axis=0)  # [3]
+        out_dict["shape"][i] = torch.from_numpy(object_shape)
+
+        valid_steps = np.where(valid)[0]
+        position = states[:, :3]  # [n_step, dim], x, y, z
+        velocity = states[:, 7:9]  # [n_step, 2], vx, vy
+        heading = states[:, 6]  # [n_step], heading
+
+        # valid.sum() should > 1:
+        t_start, t_end = valid_steps[0], valid_steps[-1]
+        f_pos = interp1d(valid_steps, position[valid], axis=0)
+        f_vel = interp1d(valid_steps, velocity[valid], axis=0)
+        f_yaw = interp1d(valid_steps, np.unwrap(heading[valid], axis=0), axis=0)
+        t_in = np.arange(t_start, t_end + 1)
+        out_dict["valid_mask"][i, t_start : t_end + 1] = True
+        out_dict["position"][i, t_start : t_end + 1] = torch.from_numpy(f_pos(t_in))
+        out_dict["velocity"][i, t_start : t_end + 1] = torch.from_numpy(f_vel(t_in))
+        out_dict["heading"][i, t_start : t_end + 1] = torch.from_numpy(f_yaw(t_in))
+
+    out_dict['av_idx'] = out_dict['id'].tolist().index(av_id)
+
+    return out_dict
+
+
+def get_map_features(map_infos, tf_current_light, dim=3):
+    lane_segments = map_infos['lane']
+    all_polylines = map_infos["all_polylines"]
+    crosswalks = map_infos['crosswalk']
+    road_edges = map_infos['road_edge']
+    road_lines = map_infos['road_line']
+    lane_segment_ids = [info["id"] for info in lane_segments]
+    cross_walk_ids = [info["id"] for info in crosswalks]
+    road_edge_ids = [info["id"] for info in road_edges]
+    road_line_ids = [info["id"] for info in road_lines]
+    polygon_ids = lane_segment_ids + road_edge_ids + road_line_ids + cross_walk_ids
+    num_polygons = len(lane_segment_ids) + len(road_edge_ids) + len(road_line_ids) + len(cross_walk_ids)
+
+    # initialization
+    polygon_type = torch.zeros(num_polygons, dtype=torch.uint8)
+    polygon_light_type = torch.ones(num_polygons, dtype=torch.uint8) * 3
+
+    # list of (num_of_segments,), each element has shape of (num_of_points_of_current_segment - 1, dim)
+    point_position: List[Optional[torch.Tensor]] = [None] * num_polygons
+    point_orientation: List[Optional[torch.Tensor]] = [None] * num_polygons
+    point_magnitude: List[Optional[torch.Tensor]] = [None] * num_polygons
+    point_height: List[Optional[torch.Tensor]] = [None] * num_polygons
+    point_type: List[Optional[torch.Tensor]] = [None] * num_polygons
+
+    for lane_segment in lane_segments:
+        lane_segment = easydict.EasyDict(lane_segment)
+        lane_segment_idx = polygon_ids.index(lane_segment.id)
+        polyline_index = lane_segment.polyline_index # (start index of point in current scenario, end index of point in current scenario)
+        centerline = all_polylines[polyline_index[0] : polyline_index[1], :] # （num_of_points_of_current_segment, 5)
+        centerline = torch.from_numpy(centerline).float()
+        polygon_type[lane_segment_idx] = _polygon_types.index(Lane_type_hash[lane_segment.type])
+
+        res = tf_current_light[tf_current_light["lane_id"] == str(lane_segment.id)]
+        if len(res) != 0:
+            polygon_light_type[lane_segment_idx] = _polygon_light_type.index(res["state"].item())
+
+        point_position[lane_segment_idx] = torch.cat([centerline[:-1, :dim]], dim=0) # (num_of_points_of_current_segment - 1, 3)
+        center_vectors = centerline[1:] - centerline[:-1] # (num_of_points_of_current_segment - 1, 5)
+        point_orientation[lane_segment_idx] = torch.cat([torch.atan2(center_vectors[:, 1], center_vectors[:, 0])], dim=0) # (num_of_points_of_current_segment - 1,)
+        point_magnitude[lane_segment_idx] = torch.norm(torch.cat([center_vectors[:, :2]], dim=0), p=2, dim=-1) # (num_of_points_of_current_segment - 1,)
+        point_height[lane_segment_idx] = torch.cat([center_vectors[:, 2]], dim=0) # (num_of_points_of_current_segment - 1,)
+        center_type = _point_types.index('CENTERLINE')
+        point_type[lane_segment_idx] = torch.cat(
+            [torch.full((len(center_vectors),), center_type, dtype=torch.uint8)], dim=0)
+
+    for lane_segment in road_edges:
+        lane_segment = easydict.EasyDict(lane_segment)
+        lane_segment_idx = polygon_ids.index(lane_segment.id)
+        polyline_index = lane_segment.polyline_index
+        centerline = all_polylines[polyline_index[0] : polyline_index[1], :]
+        centerline = torch.from_numpy(centerline).float()
+        polygon_type[lane_segment_idx] = _polygon_types.index("VEHICLE")
+
+        point_position[lane_segment_idx] = torch.cat([centerline[:-1, :dim]], dim=0)
+        center_vectors = centerline[1:] - centerline[:-1]
+        point_orientation[lane_segment_idx] = torch.cat([torch.atan2(center_vectors[:, 1], center_vectors[:, 0])], dim=0)
+        point_magnitude[lane_segment_idx] = torch.norm(torch.cat([center_vectors[:, :2]], dim=0), p=2, dim=-1)
+        point_height[lane_segment_idx] = torch.cat([center_vectors[:, 2]], dim=0)
+        center_type = _point_types.index('EDGE')
+        point_type[lane_segment_idx] = torch.cat(
+            [torch.full((len(center_vectors),), center_type, dtype=torch.uint8)], dim=0)
+
+    for lane_segment in road_lines:
+        lane_segment = easydict.EasyDict(lane_segment)
+        lane_segment_idx = polygon_ids.index(lane_segment.id)
+        polyline_index = lane_segment.polyline_index
+        centerline = all_polylines[polyline_index[0] : polyline_index[1], :]
+        centerline = torch.from_numpy(centerline).float()
+
+        polygon_type[lane_segment_idx] = _polygon_types.index("VEHICLE")
+
+        point_position[lane_segment_idx] = torch.cat([centerline[:-1, :dim]], dim=0)
+        center_vectors = centerline[1:] - centerline[:-1]
+        point_orientation[lane_segment_idx] = torch.cat([torch.atan2(center_vectors[:, 1], center_vectors[:, 0])], dim=0)
+        point_magnitude[lane_segment_idx] = torch.norm(torch.cat([center_vectors[:, :2]], dim=0), p=2, dim=-1)
+        point_height[lane_segment_idx] = torch.cat([center_vectors[:, 2]], dim=0)
+        center_type = _point_types.index(boundary_type_hash[lane_segment.type])
+        point_type[lane_segment_idx] = torch.cat(
+            [torch.full((len(center_vectors),), center_type, dtype=torch.uint8)], dim=0)
+
+    for crosswalk in crosswalks:
+        crosswalk = easydict.EasyDict(crosswalk)
+        lane_segment_idx = polygon_ids.index(crosswalk.id)
+        polyline_index = crosswalk.polyline_index
+        centerline = all_polylines[polyline_index[0] : polyline_index[1], :]
+        centerline = torch.from_numpy(centerline).float()
+
+        polygon_type[lane_segment_idx] = _polygon_types.index("PEDESTRIAN")
+
+        point_position[lane_segment_idx] = torch.cat([centerline[:-1, :dim]], dim=0)
+        center_vectors = centerline[1:] - centerline[:-1]
+        point_orientation[lane_segment_idx] = torch.cat([torch.atan2(center_vectors[:, 1], center_vectors[:, 0])], dim=0)
+        point_magnitude[lane_segment_idx] = torch.norm(torch.cat([center_vectors[:, :2]], dim=0), p=2, dim=-1)
+        point_height[lane_segment_idx] = torch.cat([center_vectors[:, 2]], dim=0)
+        center_type = _point_types.index("CROSSWALK")
+        point_type[lane_segment_idx] = torch.cat(
+            [torch.full((len(center_vectors),), center_type, dtype=torch.uint8)], dim=0)
+
+    # (num_of_segments,), each element represents the number of points of the segment
+    num_points = torch.tensor([point.size(0) for point in point_position], dtype=torch.long)
+    # (2, total_num_of_points_of_all_segments), store the point index of segment and its corresponding segment index
+    # e.g. a scenario has 203 segments, and totally 14039 points:
+    # tensor([[    0,     1,     2,  ..., 14927, 14928, 14929],
+    #         [    0,     0,     0,  ...,   202,   202,   202]]) => polygon_ids.index(lane_segment.id)
+    point_to_polygon_edge_index = torch.stack(
+        [torch.arange(num_points.sum(), dtype=torch.long),
+            torch.arange(num_polygons, dtype=torch.long).repeat_interleave(num_points)], dim=0)
+    # list of (num_of_lane_segments,)
+    polygon_to_polygon_edge_index = []
+    # list of (num_of_lane_segments,)
+    polygon_to_polygon_type = []
+    for lane_segment in lane_segments:
+        lane_segment = easydict.EasyDict(lane_segment)
+        lane_segment_idx = polygon_ids.index(lane_segment.id)
+        pred_inds = []
+        for pred in lane_segment.entry_lanes:
+            pred_idx = safe_list_index(polygon_ids, pred)
+            if pred_idx is not None:
+                pred_inds.append(pred_idx)
+        if len(pred_inds) != 0:
+            polygon_to_polygon_edge_index.append(
+                torch.stack([torch.tensor(pred_inds, dtype=torch.long),
+                             torch.full((len(pred_inds),), lane_segment_idx, dtype=torch.long)], dim=0))
+            polygon_to_polygon_type.append(
+                torch.full((len(pred_inds),), _polygon_to_polygon_types.index('PRED'), dtype=torch.uint8))
+        succ_inds = []
+        for succ in lane_segment.exit_lanes:
+            succ_idx = safe_list_index(polygon_ids, succ)
+            if succ_idx is not None:
+                succ_inds.append(succ_idx)
+        if len(succ_inds) != 0:
+            polygon_to_polygon_edge_index.append(
+                torch.stack([torch.tensor(succ_inds, dtype=torch.long),
+                             torch.full((len(succ_inds),), lane_segment_idx, dtype=torch.long)], dim=0))
+            polygon_to_polygon_type.append(
+                torch.full((len(succ_inds),), _polygon_to_polygon_types.index('SUCC'), dtype=torch.uint8))
+        if len(lane_segment.left_neighbors) != 0:
+            left_neighbor_ids = lane_segment.left_neighbors
+            for left_neighbor_id in left_neighbor_ids:
+                left_idx = safe_list_index(polygon_ids, left_neighbor_id)
+                if left_idx is not None:
+                    polygon_to_polygon_edge_index.append(
+                        torch.tensor([[left_idx], [lane_segment_idx]], dtype=torch.long))
+                    polygon_to_polygon_type.append(
+                        torch.tensor([_polygon_to_polygon_types.index('LEFT')], dtype=torch.uint8))
+        if len(lane_segment.right_neighbors) != 0:
+            right_neighbor_ids = lane_segment.right_neighbors
+            for right_neighbor_id in right_neighbor_ids:
+                right_idx = safe_list_index(polygon_ids, right_neighbor_id)
+                if right_idx is not None:
+                    polygon_to_polygon_edge_index.append(
+                        torch.tensor([[right_idx], [lane_segment_idx]], dtype=torch.long))
+                    polygon_to_polygon_type.append(
+                        torch.tensor([_polygon_to_polygon_types.index('RIGHT')], dtype=torch.uint8))
+    if len(polygon_to_polygon_edge_index) != 0:
+        polygon_to_polygon_edge_index = torch.cat(polygon_to_polygon_edge_index, dim=1)
+        polygon_to_polygon_type = torch.cat(polygon_to_polygon_type, dim=0)
+    else:
+        polygon_to_polygon_edge_index = torch.tensor([[], []], dtype=torch.long)
+        polygon_to_polygon_type = torch.tensor([], dtype=torch.uint8)
+
+    map_data = {
+        'map_polygon': {},
+        'map_point': {},
+        ('map_point', 'to', 'map_polygon'): {},
+        ('map_polygon', 'to', 'map_polygon'): {},
+    }
+    map_data['map_polygon']['num_nodes'] = num_polygons # int, number of map segments in the scenario
+    map_data['map_polygon']['type'] = polygon_type # (num_polygons,) type of each polygon
+    map_data['map_polygon']['light_type'] = polygon_light_type # (num_polygons,) light type of each polygon, 3 means unknown
+    if len(num_points) == 0:
+        map_data['map_point']['num_nodes'] = 0
+        map_data['map_point']['position'] = torch.tensor([], dtype=torch.float)
+        map_data['map_point']['orientation'] = torch.tensor([], dtype=torch.float)
+        map_data['map_point']['magnitude'] = torch.tensor([], dtype=torch.float)
+        if dim == 3:
+            map_data['map_point']['height'] = torch.tensor([], dtype=torch.float)
+        map_data['map_point']['type'] = torch.tensor([], dtype=torch.uint8)
+        map_data['map_point']['side'] = torch.tensor([], dtype=torch.uint8)
+    else:
+        map_data['map_point']['num_nodes'] = num_points.sum().item() # int, number of total points of all segments in the scenario
+        map_data['map_point']['position'] = torch.cat(point_position, dim=0) # (num_of_total_points_of_all_segments, 3)
+        map_data['map_point']['orientation'] = torch.cat(point_orientation, dim=0) # (num_of_total_points_of_all_segments,)
+        map_data['map_point']['magnitude'] = torch.cat(point_magnitude, dim=0) # (num_of_total_points_of_all_segments,)
+        if dim == 3:
+            map_data['map_point']['height'] = torch.cat(point_height, dim=0) # (num_of_total_points_of_all_segments,)
+        map_data['map_point']['type'] = torch.cat(point_type, dim=0) # (num_of_total_points_of_all_segments,) type of point => `_point_types`
+    map_data['map_point', 'to', 'map_polygon']['edge_index'] = point_to_polygon_edge_index # (2, num_of_total_points_of_all_segments)
+    map_data['map_polygon', 'to', 'map_polygon']['edge_index'] = polygon_to_polygon_edge_index
+    map_data['map_polygon', 'to', 'map_polygon']['type'] = polygon_to_polygon_type
+
+    if int(os.getenv('DEBUG_MAP', 1)):
+        import matplotlib.pyplot as plt
+        plt.axis('equal')
+        plt.scatter(map_data['map_point']['position'][:, 0],
+                    map_data['map_point']['position'][:, 1], s=0.2, c='black', edgecolors='none')
+        plt.savefig("debug.png", dpi=600)
+
+    return map_data
+
+
+# def process_agent(track_info, tracks_to_predict, scenario_id, start_timestamp, end_timestamp):
+
+#     agents_array = track_info["states"].transpose(1, 0, 2) # (num_timesteps, num_agents, 10) e.g. (91, 15, 10)
+#     object_id = np.array(track_info["object_id"]) # (num_agents,) global id of each agent
+#     object_type = track_info["object_type"] # (num_agents,) type of each agent, e.g. 'TYPE_VEHICLE'
+#     id_hash = {object_id[o_idx]: object_type[o_idx] for o_idx in range(len(object_id))}
+
+#     def type_hash(x):
+#         tp = id_hash[x]
+#         type_re_hash = {
+#             "TYPE_VEHICLE": "vehicle",
+#             "TYPE_PEDESTRIAN": "pedestrian",
+#             "TYPE_CYCLIST": "cyclist",
+#             "TYPE_OTHER": "background",
+#             "TYPE_UNSET": "background"
+#         }
+#         return type_re_hash[tp]
+
+#     columns = ['observed', 'track_id', 'object_type', 'object_category', 'timestep',
+#                'position_x', 'position_y', 'position_z', 'length', 'width', 'height', 'heading', 'velocity_x', 'velocity_y',
+#                'scenario_id', 'start_timestamp', 'end_timestamp', 'num_timestamps',
+#                'focal_track_id', 'city', 'validity']
+
+#     # (num_timesteps, num_agents, 10) e.g. (91, 15, 10)
+#     new_columns = np.ones((agents_array.shape[0], agents_array.shape[1], 11))
+#     new_columns[:11, :, 0] = True # observed, 10 timesteps
+#     new_columns[11:, :, 0] = False # not observed (current + future) 
+#     for index in range(new_columns.shape[0]):
+#         new_columns[index, :, 4] = int(index) # timestep (0 ~ 90)
+#     new_columns[..., 1] = object_id
+#     new_columns[..., 2] = object_id
+#     new_columns[:, tracks_to_predict['track_index'], 3] = 3
+#     new_columns[..., 5] = 11
+#     new_columns[..., 6] = int(start_timestamp) # 0
+#     new_columns[..., 7] = int(end_timestamp)   # 91
+#     new_columns[..., 8] = int(91)              # 91
+#     new_columns[..., 9] = object_id
+#     new_columns[..., 10] = 10086
+#     new_columns = new_columns
+#     new_agents_array = np.concatenate([new_columns, agents_array], axis=-1) # (num_timesteps, num_agents, 21) e.g. (91, 15, 21)
+#     # filter out the invalid timestep of agents, reshape to (num_valid_of_timesteps_of_all_agents, 21) e.g. (91, 15, 21) -> (1137, 21)
+#     if args.disable_invalid:
+#         new_agents_array = new_agents_array[new_agents_array[..., -1] == 1.0].reshape(-1, new_agents_array.shape[-1])
+#     else:
+#         agent_valid_mask = new_agents_array[..., -1] # (num_timesteps, num_agents)
+#         agent_mask = np.sum(agent_valid_mask, axis=0) > MIN_VALID_STEPS # NOTE: 10 is a empirical parameter
+#         new_agents_array = new_agents_array[:, agent_mask]
+#         new_agents_array = new_agents_array.reshape(-1, new_agents_array.shape[-1]) # (91, 15, 21) -> (1365, 21)
+#     new_agents_array = new_agents_array[..., [0, 1, 2, 3, 4, 11, 12, 13, 14, 15, 16, 17, 18, 19, 5, 6, 7, 8, 9, 10, 20]]
+#     new_agents_array = pd.DataFrame(data=new_agents_array, columns=columns)
+#     new_agents_array["object_type"] = new_agents_array["object_type"].apply(func=type_hash)
+#     new_agents_array["start_timestamp"] = new_agents_array["start_timestamp"].astype(int)
+#     new_agents_array["end_timestamp"] = new_agents_array["end_timestamp"].astype(int)
+#     new_agents_array["num_timestamps"] = new_agents_array["num_timestamps"].astype(int)
+#     new_agents_array["scenario_id"] = scenario_id
+
+#     return new_agents_array
+
+
+def process_dynamic_map(dynamic_map_infos):
+    lane_ids = dynamic_map_infos["lane_id"]
+    tf_lights = []
+    for t in range(len(lane_ids)):
+        lane_id = lane_ids[t]
+        time = np.ones_like(lane_id) * t
+        state = dynamic_map_infos["state"][t]
+        tf_light = np.concatenate([lane_id, time, state], axis=0)
+        tf_lights.append(tf_light)
+    tf_lights = np.concatenate(tf_lights, axis=1).transpose(1, 0)
+    tf_lights = pd.DataFrame(data=tf_lights, columns=["lane_id", "time_step", "state"])
+    tf_lights["time_step"] = tf_lights["time_step"].astype("str")
+    tf_lights["lane_id"] = tf_lights["lane_id"].astype("str")
+    tf_lights["state"] = tf_lights["state"].astype("str")
+    tf_lights.loc[tf_lights["state"].str.contains("STOP"), ["state"]] = (
+        "LANE_STATE_STOP"
+    )
+    tf_lights.loc[tf_lights["state"].str.contains("GO"), ["state"]] = "LANE_STATE_GO"
+    tf_lights.loc[tf_lights["state"].str.contains("CAUTION"), ["state"]] = (
+        "LANE_STATE_CAUTION"
+    )
+    tf_lights.loc[tf_lights["state"].str.contains("UNKNOWN"), ["state"]] = (
+        "LANE_STATE_UNKNOWN"
+    )
+
+    return tf_lights
+
+
+polyline_type = {
+    # for lane
+    'TYPE_UNDEFINED': -1,
+    'TYPE_FREEWAY': 1,
+    'TYPE_SURFACE_STREET': 2,
+    'TYPE_BIKE_LANE': 3,
+
+    # for roadline
+    'TYPE_UNKNOWN': -1,
+    'TYPE_BROKEN_SINGLE_WHITE': 6,
+    'TYPE_SOLID_SINGLE_WHITE': 7,
+    'TYPE_SOLID_DOUBLE_WHITE': 8,
+    'TYPE_BROKEN_SINGLE_YELLOW': 9,
+    'TYPE_BROKEN_DOUBLE_YELLOW': 10,
+    'TYPE_SOLID_SINGLE_YELLOW': 11,
+    'TYPE_SOLID_DOUBLE_YELLOW': 12,
+    'TYPE_PASSING_DOUBLE_YELLOW': 13,
+
+    # for roadedge
+    'TYPE_ROAD_EDGE_BOUNDARY': 15,
+    'TYPE_ROAD_EDGE_MEDIAN': 16,
+
+    # for stopsign
+    'TYPE_STOP_SIGN': 17,
+
+    # for crosswalk
+    'TYPE_CROSSWALK': 18,
+
+    # for speed bump
+    'TYPE_SPEED_BUMP': 19
+}
+
+object_type = {
+    0: 'TYPE_UNSET',
+    1: 'TYPE_VEHICLE',
+    2: 'TYPE_PEDESTRIAN',
+    3: 'TYPE_CYCLIST',
+    4: 'TYPE_OTHER'
+}
+
+
+def decode_tracks_from_proto(scenario):
+    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]
+
+    track_infos = {
+        'object_id': [],  # {0: unset, 1: vehicle, 2: pedestrian, 3: cyclist, 4: others}
+        'object_type': [],
+        'states': [],
+        'valid': [],
+        'role': [],
+    }
+
+    # tracks mean N number of objects, e.g. len(tracks) = 55
+    # each track has 91 states, e.g. len(tracks[0].states) == 91
+    # each state has 10 attributes: center_x, center_y, center_z, length, ..., velocity_y, valid
+    for i, cur_data in enumerate(scenario.tracks):
+
+        step_state = []
+        step_valid = []
+
+        for s in cur_data.states:  # n_steps
+            step_state.append(
+                [
+                    s.center_x,
+                    s.center_y,
+                    s.center_z,
+                    s.length,
+                    s.width,
+                    s.height,
+                    s.heading,
+                    s.velocity_x,
+                    s.velocity_y,
+                ]
+            )
+            step_valid.append(s.valid)
+            # This angle is normalized to [-pi, pi). The velocity vector in m/s
+
+        track_infos['object_id'].append(cur_data.id) # id of object in this track
+        track_infos['object_type'].append(cur_data.object_type - 1)
+        track_infos['states'].append(np.array(step_state, dtype=np.float32))
+        track_infos['valid'].append(np.array(step_valid))
+
+        track_infos['role'].append([False, False, False])
+        if i in track_index_predict:
+            track_infos['role'][-1][2] = True  # predict=2
+        if cur_data.id in object_id_interest:
+            track_infos['role'][-1][1] = True  # interest=1
+        if i == sdc_track_index:
+            track_infos['role'][-1][0] = True  # ego_vehicle=0
+
+    track_infos['states']   = np.array(track_infos['states'], dtype=np.float32)  # (n_agent, n_step, 9)
+    track_infos['valid']    = np.array(track_infos['valid'], dtype=np.bool_)
+    track_infos['role']     = np.array(track_infos['role'], dtype=np.bool_)
+    track_infos['object_id'] = np.array(track_infos['object_id'], dtype=np.int64)
+    track_infos['object_type'] = np.array(track_infos['object_type'], dtype=np.uint8)
+    track_infos['tracks_to_predict'] = np.array(track_index_predict, dtype=np.int64)
+
+    return track_infos
+
+
+from collections import defaultdict
+
+def decode_map_features_from_proto(map_features):
+    map_infos = {
+        'lane': [],
+        'road_line': [],
+        'road_edge': [],
+        'stop_sign': [],
+        'crosswalk': [],
+        'speed_bump': [],
+        'lane_dict': {},
+        'lane2other_dict': {}
+    }
+    polylines = []
+
+    point_cnt = 0
+    lane2other_dict = defaultdict(list)
+
+    for cur_data in map_features:
+        cur_info = {'id': cur_data.id}
+
+        if cur_data.lane.ByteSize() > 0:
+            cur_info['speed_limit_mph'] = cur_data.lane.speed_limit_mph
+            cur_info['type'] = cur_data.lane.type + 1  # 0: undefined, 1: freeway, 2: surface_street, 3: bike_lane
+            cur_info['left_neighbors'] = [lane.feature_id for lane in cur_data.lane.left_neighbors]
+
+            cur_info['right_neighbors'] = [lane.feature_id for lane in cur_data.lane.right_neighbors]
+
+            cur_info['interpolating'] = cur_data.lane.interpolating
+            cur_info['entry_lanes'] = list(cur_data.lane.entry_lanes)
+            cur_info['exit_lanes'] = list(cur_data.lane.exit_lanes)
+
+            cur_info['left_boundary_type'] = [x.boundary_type + 5 for x in cur_data.lane.left_boundaries]
+            cur_info['right_boundary_type'] = [x.boundary_type + 5 for x in cur_data.lane.right_boundaries]
+
+            cur_info['left_boundary'] = [x.boundary_feature_id for x in cur_data.lane.left_boundaries]
+            cur_info['right_boundary'] = [x.boundary_feature_id for x in cur_data.lane.right_boundaries]
+            cur_info['left_boundary_start_index'] = [lane.lane_start_index for lane in cur_data.lane.left_boundaries]
+            cur_info['left_boundary_end_index'] = [lane.lane_end_index for lane in cur_data.lane.left_boundaries]
+            cur_info['right_boundary_start_index'] = [lane.lane_start_index for lane in cur_data.lane.right_boundaries]
+            cur_info['right_boundary_end_index'] = [lane.lane_end_index for lane in cur_data.lane.right_boundaries]
+
+            lane2other_dict[cur_data.id].extend(cur_info['left_boundary'])
+            lane2other_dict[cur_data.id].extend(cur_info['right_boundary'])
+
+            global_type = cur_info['type']
+            cur_polyline = np.stack(
+                [np.array([point.x, point.y, point.z, global_type, cur_data.id]) for point in cur_data.lane.polyline],
+                axis=0)
+            cur_polyline = np.concatenate((cur_polyline[:, 0:3], cur_polyline[:, 3:]), axis=-1)
+            if cur_polyline.shape[0] <= 1:
+                continue
+            map_infos['lane'].append(cur_info)
+            map_infos['lane_dict'][cur_data.id] = cur_info
+
+        elif cur_data.road_line.ByteSize() > 0:
+            cur_info['type'] = cur_data.road_line.type + 5
+
+            global_type = cur_info['type']
+            cur_polyline = np.stack([np.array([point.x, point.y, point.z, global_type, cur_data.id]) for point in
+                                     cur_data.road_line.polyline], axis=0)
+            cur_polyline = np.concatenate((cur_polyline[:, 0:3], cur_polyline[:, 3:]), axis=-1)
+            if cur_polyline.shape[0] <= 1:
+                continue
+            map_infos['road_line'].append(cur_info) # (num_points, 5)
+
+        elif cur_data.road_edge.ByteSize() > 0:
+            cur_info['type'] = cur_data.road_edge.type + 14
+
+            global_type = cur_info['type']
+            cur_polyline = np.stack([np.array([point.x, point.y, point.z, global_type, cur_data.id]) for point in
+                                     cur_data.road_edge.polyline], axis=0)
+            cur_polyline = np.concatenate((cur_polyline[:, 0:3], cur_polyline[:, 3:]), axis=-1)
+            if cur_polyline.shape[0] <= 1:
+                continue
+            map_infos['road_edge'].append(cur_info)
+
+        elif cur_data.stop_sign.ByteSize() > 0:
+            cur_info['lane_ids'] = list(cur_data.stop_sign.lane)
+            for i in cur_info['lane_ids']:
+                lane2other_dict[i].append(cur_data.id)
+            point = cur_data.stop_sign.position
+            cur_info['position'] = np.array([point.x, point.y, point.z])
+
+            global_type = polyline_type['TYPE_STOP_SIGN']
+            cur_polyline = np.array([point.x, point.y, point.z, global_type, cur_data.id]).reshape(1, 5)
+            if cur_polyline.shape[0] <= 1:
+                continue
+            map_infos['stop_sign'].append(cur_info)
+        elif cur_data.crosswalk.ByteSize() > 0:
+            global_type = polyline_type['TYPE_CROSSWALK']
+            cur_polyline = np.stack([np.array([point.x, point.y, point.z, global_type, cur_data.id]) for point in
+                                     cur_data.crosswalk.polygon], axis=0)
+            cur_polyline = np.concatenate((cur_polyline[:, 0:3], cur_polyline[:, 3:]), axis=-1)
+            if cur_polyline.shape[0] <= 1:
+                continue
+            map_infos['crosswalk'].append(cur_info)
+
+        elif cur_data.speed_bump.ByteSize() > 0:
+            global_type = polyline_type['TYPE_SPEED_BUMP']
+            cur_polyline = np.stack([np.array([point.x, point.y, point.z, global_type, cur_data.id]) for point in
+                                     cur_data.speed_bump.polygon], axis=0)
+            cur_polyline = np.concatenate((cur_polyline[:, 0:3], cur_polyline[:, 3:]), axis=-1)
+            if cur_polyline.shape[0] <= 1:
+                continue
+            map_infos['speed_bump'].append(cur_info)
+
+        else:
+            continue
+        polylines.append(cur_polyline)
+        cur_info['polyline_index'] = (point_cnt, point_cnt + len(cur_polyline)) # (start index of point in current scenario, end index of point in current scenario)
+        point_cnt += len(cur_polyline)
+
+    polylines = np.concatenate(polylines, axis=0).astype(np.float32)
+    map_infos['all_polylines'] = polylines # (num_of_total_points_in_current_scenario, 5)
+    map_infos['lane2other_dict'] = lane2other_dict
+    return map_infos
+
+
+def decode_dynamic_map_states_from_proto(dynamic_map_states):
+
+    signal_state = {
+        0: 'LANE_STATE_UNKNOWN',
+        # States for traffic signals with arrows.
+        1: 'LANE_STATE_ARROW_STOP',
+        2: 'LANE_STATE_ARROW_CAUTION',
+        3: 'LANE_STATE_ARROW_GO',
+        # Standard round traffic signals.
+        4: 'LANE_STATE_STOP',
+        5: 'LANE_STATE_CAUTION',
+        6: 'LANE_STATE_GO',
+        # Flashing light signals.
+        7: 'LANE_STATE_FLASHING_STOP',
+        8: 'LANE_STATE_FLASHING_CAUTION'
+    }
+
+    dynamic_map_infos = {
+        'lane_id': [],
+        'state': [],
+        'stop_point': []
+    }
+    for cur_data in dynamic_map_states:  # len(dynamic_map_states) = num_timestamp
+        lane_id, state, stop_point = [], [], []
+        for cur_signal in cur_data.lane_states:  # (num_observed_signals)
+            lane_id.append(cur_signal.lane)
+            state.append(signal_state[cur_signal.state])
+            stop_point.append([cur_signal.stop_point.x, cur_signal.stop_point.y, cur_signal.stop_point.z])
+
+        dynamic_map_infos['lane_id'].append(np.array([lane_id]))
+        dynamic_map_infos['state'].append(np.array([state]))
+        dynamic_map_infos['stop_point'].append(np.array([stop_point]))
+
+    return dynamic_map_infos
+
+
+# def process_single_data(scenario):
+#     info = {}
+#     info['scenario_id'] = scenario.scenario_id
+#     info['timestamps_seconds'] = list(scenario.timestamps_seconds)  # list of int of shape (91)
+#     info['current_time_index'] = scenario.current_time_index  # int, 10
+#     info['sdc_track_index'] = scenario.sdc_track_index  # int
+#     info['objects_of_interest'] = list(scenario.objects_of_interest)  # list, could be empty list
+
+#     info['tracks_to_predict'] = {
+#         'track_index': [cur_pred.track_index for cur_pred in scenario.tracks_to_predict],
+#         'difficulty': [cur_pred.difficulty for cur_pred in scenario.tracks_to_predict]
+#     }  # for training: suggestion of objects to train on, for val/test: need to be predicted
+
+#     # decode tracks data
+#     track_infos = decode_tracks_from_proto(scenario.tracks)
+#     info['tracks_to_predict']['object_type'] = [track_infos['object_type'][cur_idx] for cur_idx in
+#                                                 info['tracks_to_predict']['track_index']]
+#     # decode map related data
+#     map_infos = decode_map_features_from_proto(scenario.map_features)
+#     dynamic_map_infos = decode_dynamic_map_states_from_proto(scenario.dynamic_map_states)
+
+#     save_infos = {
+#         'track_infos': track_infos,
+#         'map_infos': map_infos,
+#         'dynamic_map_infos': dynamic_map_infos,
+#     }
+#     save_infos.update(info)
+#     return save_infos
+
+
+def wm2argo(file, input_dir, output_dir, existing_files=[], output_dir_tfrecords_splitted=None):
+    file_path = os.path.join(input_dir, file)
+    dataset = tf.data.TFRecordDataset(file_path, compression_type='', num_parallel_reads=3)
+
+    for cnt, tf_data in tqdm(enumerate(dataset), leave=False, desc=f'Process {file}...'):
+
+        scenario = scenario_pb2.Scenario()
+        scenario.ParseFromString(bytearray(tf_data.numpy()))
+        scenario_id = scenario.scenario_id
+        tqdm.write(f"idx: {cnt}, scenario_id: {scenario_id} of {file}")
+
+        if f'{scenario_id}.pkl' not in existing_files:
+
+            map_infos = decode_map_features_from_proto(scenario.map_features)
+            track_infos = decode_tracks_from_proto(scenario)
+            dynamic_map_infos = decode_dynamic_map_states_from_proto(scenario.dynamic_map_states)
+            sdc_track_index = scenario.sdc_track_index  # int
+            av_id = track_infos['object_id'][sdc_track_index]
+            # if len(track_infos['tracks_to_predict']) < 1:
+            #     return
+
+            current_time_index = scenario.current_time_index
+            tf_lights = process_dynamic_map(dynamic_map_infos)
+            tf_current_light = tf_lights.loc[tf_lights["time_step"] == current_time_index] # 10 (history) + 1 (current) + 80 (future)
+            map_data = get_map_features(map_infos, tf_current_light)
+
+            # new_agents_array = process_agent(track_infos, tracks_to_predict, scenario_id, 0, 91) # mtr2argo
+            data = dict()
+            data.update(map_data)
+            data['scenario_id'] = scenario_id
+            data['agent'] = get_agent_features(track_infos, av_id, num_historical_steps=current_time_index + 1, num_steps=91)
+
+            with open(os.path.join(output_dir, f'{scenario_id}.pkl'), "wb+") as f:
+                pickle.dump(data, f)
+
+        if output_dir_tfrecords_splitted is not None:
+            tf_file = os.path.join(output_dir_tfrecords_splitted, f'{scenario_id}.tfrecords')
+            if not os.path.exists(tf_file):
+                with tf.io.TFRecordWriter(tf_file) as file_writer:
+                    file_writer.write(tf_data.numpy())
+
+
+def batch_process9s_transformer(input_dir, output_dir, split, num_workers=2):
+    signal.signal(signal.SIGINT, signal.SIG_IGN)
+
+    output_dir_tfrecords_splitted = None
+    if split == "validation":
+        output_dir_tfrecords_splitted = os.path.join(output_dir, 'validation_tfrecords_splitted')
+        os.makedirs(output_dir_tfrecords_splitted, exist_ok=True)
+
+    input_dir = os.path.join(input_dir, split)
+    output_dir = os.path.join(output_dir, split)
+    os.makedirs(output_dir, exist_ok=True)
+
+    packages = sorted(os.listdir(input_dir))
+    existing_files = sorted(os.listdir(output_dir))
+    func = partial(
+        wm2argo,
+        output_dir=output_dir,
+        input_dir=input_dir,
+        existing_files=existing_files,
+        output_dir_tfrecords_splitted=output_dir_tfrecords_splitted
+    )
+    try:
+        with multiprocessing.Pool(num_workers, maxtasksperchild=10) as p:
+            r = list(tqdm(p.imap_unordered(func, packages), total=len(packages)))
+    except KeyboardInterrupt:
+        p.terminate()
+        p.join()
+
+
+def generate_meta_infos(data_dir):
+    import json
+
+    meta_infos = dict()
+
+    for split in tqdm(['training', 'validation', 'test'], leave=False):
+        if not os.path.exists(os.path.join(data_dir, split)):
+            continue
+
+        split_infos = dict()
+        files = os.listdir(os.path.join(data_dir, split))
+        for file in tqdm(files, leave=False):
+            try:
+                data = pickle.load(open(os.path.join(data_dir, split, file), 'rb'))
+            except Exception as e:
+                tqdm.write(f'Failed to load scenario {file} due to {e}')
+                continue
+            scenario_infos = dict(num_agents=data['agent']['num_nodes'])
+            scenario_id = data['scenario_id']
+            split_infos[scenario_id] = scenario_infos
+
+        meta_infos[split] = split_infos
+
+    with open(os.path.join(data_dir, 'meta_infos.json'), 'w', encoding='utf-8') as f:
+        json.dump(meta_infos, f, indent=4)
+
+
+if __name__ == "__main__":
+    parser = ArgumentParser()
+    parser.add_argument('--input_dir', type=str, default='data/waymo/')
+    parser.add_argument('--output_dir', type=str, default='data/waymo_processed/')
+    parser.add_argument('--split', type=str, default='validation')
+    parser.add_argument('--no_batch', action='store_true')
+    parser.add_argument('--disable_invalid', action="store_true")
+    parser.add_argument('--generate_meta_infos', action="store_true")
+    args = parser.parse_args()
+
+    if args.generate_meta_infos:
+        generate_meta_infos(args.output_dir)
+
+    elif args.no_batch:
+
+        output_dir_tfrecords_splitted = None
+        if args.split == "validation":
+            output_dir_tfrecords_splitted = os.path.join(args.output_dir, 'validation_tfrecords_splitted')
+            os.makedirs(output_dir_tfrecords_splitted, exist_ok=True)
+
+        input_dir = os.path.join(args.input_dir, args.split)
+        output_dir = os.path.join(args.output_dir, args.split)
+        os.makedirs(output_dir, exist_ok=True)
+
+        files = sorted(os.listdir(input_dir))
+        os.makedirs(args.output_dir, exist_ok=True)
+        for file in tqdm(files, leave=False, desc=f'Process {args.split}...'):
+            wm2argo(file, input_dir, output_dir, output_dir_tfrecords_splitted)
+
+    else:
+
+        batch_process9s_transformer(args.input_dir, args.output_dir, args.split, num_workers=96)

seq_10_150_3_3_encode_occ_separate_offsets_bs8_128_no_seqindex_long_lastvalid_ablation_grid/backups/dev/datasets/preprocess.py

@@ -0,0 +1,761 @@
+import os
+import torch
+import math
+import pickle
+import numpy as np
+from torch import nn
+from typing import Dict, Sequence
+from scipy.interpolate import interp1d
+from scipy.spatial.distance import euclidean
+from dev.utils.func import wrap_angle
+
+
+SHIFT = 5
+AGENT_SHAPE = {
+    'vehicle':      [4.3, 1.8, 1.],
+    'pedstrain':    [0.5, 0.5, 1.],
+    'cyclist':      [1.9, 0.5, 1.],
+}
+AGENT_TYPE = ['veh', 'ped', 'cyc', 'seed']
+AGENT_STATE = ['invalid', 'valid', 'enter', 'exit']
+
+
+@torch.no_grad()
+def cal_polygon_contour(pos, head, width_length) -> torch.Tensor:  # [n_agent, n_step, n_target, 4, 2]
+    x, y = pos[..., 0], pos[..., 1]  # [n_agent, n_step, n_target]
+    width, length = width_length[..., 0], width_length[..., 1]  # [n_agent, 1, 1]
+
+    half_cos = 0.5 * head.cos()  # [n_agent, n_step, n_target]
+    half_sin = 0.5 * head.sin()  # [n_agent, n_step, n_target]
+    length_cos = length * half_cos  # [n_agent, n_step, n_target]
+    length_sin = length * half_sin  # [n_agent, n_step, n_target]
+    width_cos = width * half_cos  # [n_agent, n_step, n_target]
+    width_sin = width * half_sin  # [n_agent, n_step, n_target]
+
+    left_front_x = x + length_cos - width_sin
+    left_front_y = y + length_sin + width_cos
+    left_front = torch.stack((left_front_x, left_front_y), dim=-1)
+
+    right_front_x = x + length_cos + width_sin
+    right_front_y = y + length_sin - width_cos
+    right_front = torch.stack((right_front_x, right_front_y), dim=-1)
+
+    right_back_x = x - length_cos + width_sin
+    right_back_y = y - length_sin - width_cos
+    right_back = torch.stack((right_back_x, right_back_y), dim=-1)
+
+    left_back_x = x - length_cos - width_sin
+    left_back_y = y - length_sin + width_cos
+    left_back = torch.stack((left_back_x, left_back_y), dim=-1)
+
+    polygon_contour = torch.stack(
+        (left_front, right_front, right_back, left_back), dim=-2
+    )
+
+    return polygon_contour
+
+
+def interplating_polyline(polylines, heading, distance=0.5, split_distace=5):
+    # Calculate the cumulative distance along the path, up-sample the polyline to 0.5 meter
+    dist_along_path_list = [[0]]
+    polylines_list = [[polylines[0]]]
+    for i in range(1, polylines.shape[0]):
+        euclidean_dist = euclidean(polylines[i, :2], polylines[i - 1, :2])
+        heading_diff = min(abs(max(heading[i], heading[i - 1]) - min(heading[1], heading[i - 1])),
+                           abs(max(heading[i], heading[i - 1]) - min(heading[1], heading[i - 1]) + math.pi))
+        if heading_diff > math.pi / 4 and euclidean_dist > 3:
+            dist_along_path_list.append([0])
+            polylines_list.append([polylines[i]])
+        elif heading_diff > math.pi / 8 and euclidean_dist > 3:
+            dist_along_path_list.append([0])
+            polylines_list.append([polylines[i]])
+        elif heading_diff > 0.1 and euclidean_dist > 3:
+            dist_along_path_list.append([0])
+            polylines_list.append([polylines[i]])
+        elif euclidean_dist > 10:
+            dist_along_path_list.append([0])
+            polylines_list.append([polylines[i]])
+        else:
+            dist_along_path_list[-1].append(dist_along_path_list[-1][-1] + euclidean_dist)
+            polylines_list[-1].append(polylines[i])
+    # plt.plot(polylines[:, 0], polylines[:, 1])
+    # plt.savefig('tmp.jpg')
+    new_x_list = []
+    new_y_list = []
+    multi_polylines_list = []
+    for idx in range(len(dist_along_path_list)):
+        if len(dist_along_path_list[idx]) < 2:
+            continue
+        dist_along_path = np.array(dist_along_path_list[idx])
+        polylines_cur = np.array(polylines_list[idx])
+        # Create interpolation functions for x and y coordinates
+        fx = interp1d(dist_along_path, polylines_cur[:, 0])
+        fy = interp1d(dist_along_path, polylines_cur[:, 1])
+        # fyaw = interp1d(dist_along_path, heading)
+
+        # Create an array of distances at which to interpolate
+        new_dist_along_path = np.arange(0, dist_along_path[-1], distance)
+        new_dist_along_path = np.concatenate([new_dist_along_path, dist_along_path[[-1]]])
+        # Use the interpolation functions to generate new x and y coordinates
+        new_x = fx(new_dist_along_path)
+        new_y = fy(new_dist_along_path)
+        # new_yaw = fyaw(new_dist_along_path)
+        new_x_list.append(new_x)
+        new_y_list.append(new_y)
+
+        # Combine the new x and y coordinates into a single array
+        new_polylines = np.vstack((new_x, new_y)).T
+        polyline_size = int(split_distace / distance)
+        if new_polylines.shape[0] >= (polyline_size + 1):
+            padding_size = (new_polylines.shape[0] - (polyline_size + 1)) % polyline_size
+            final_index = (new_polylines.shape[0] - (polyline_size + 1)) // polyline_size + 1
+        else:
+            padding_size = new_polylines.shape[0]
+            final_index = 0
+        multi_polylines = None
+        new_polylines = torch.from_numpy(new_polylines)
+        new_heading = torch.atan2(new_polylines[1:, 1] - new_polylines[:-1, 1],
+                                  new_polylines[1:, 0] - new_polylines[:-1, 0])
+        new_heading = torch.cat([new_heading, new_heading[-1:]], -1)[..., None]
+        new_polylines = torch.cat([new_polylines, new_heading], -1)
+        if new_polylines.shape[0] >= (polyline_size + 1):
+            multi_polylines = new_polylines.unfold(dimension=0, size=polyline_size + 1, step=polyline_size)
+            multi_polylines = multi_polylines.transpose(1, 2)
+            multi_polylines = multi_polylines[:, ::5, :]
+        if padding_size >= 3:
+            last_polyline = new_polylines[final_index * polyline_size:]
+            last_polyline = last_polyline[torch.linspace(0, last_polyline.shape[0] - 1, steps=3).long()]
+            if multi_polylines is not None:
+                multi_polylines = torch.cat([multi_polylines, last_polyline.unsqueeze(0)], dim=0)
+            else:
+                multi_polylines = last_polyline.unsqueeze(0)
+        if multi_polylines is None:
+            continue
+        multi_polylines_list.append(multi_polylines)
+    if len(multi_polylines_list) > 0:
+        multi_polylines_list = torch.cat(multi_polylines_list, dim=0)
+    else:
+        multi_polylines_list = None
+    return multi_polylines_list
+
+
+# def interplating_polyline(polylines, heading, distance=0.5, split_distance=5, device='cpu'):
+#     dist_along_path_list = [[0]]
+#     polylines_list = [[polylines[0]]]
+#     for i in range(1, polylines.shape[0]):
+#         euclidean_dist = torch.norm(polylines[i, :2] - polylines[i - 1, :2])
+#         heading_diff = min(abs(max(heading[i], heading[i - 1]) - min(heading[1], heading[i - 1])),
+#                            abs(max(heading[i], heading[i - 1]) - min(heading[1], heading[i - 1]) + torch.pi))
+#         if heading_diff > torch.pi / 4 and euclidean_dist > 3:
+#             dist_along_path_list.append([0])
+#             polylines_list.append([polylines[i]])
+#         elif heading_diff > torch.pi / 8 and euclidean_dist > 3:
+#             dist_along_path_list.append([0])
+#             polylines_list.append([polylines[i]])
+#         elif heading_diff > 0.1 and euclidean_dist > 3:
+#             dist_along_path_list.append([0])
+#             polylines_list.append([polylines[i]])
+#         elif euclidean_dist > 10:
+#             dist_along_path_list.append([0])
+#             polylines_list.append([polylines[i]])
+#         else:
+#             dist_along_path_list[-1].append(dist_along_path_list[-1][-1] + euclidean_dist)
+#             polylines_list[-1].append(polylines[i])
+
+#     new_x_list = []
+#     new_y_list = []
+#     multi_polylines_list = []
+
+#     for idx in range(len(dist_along_path_list)):
+#         if len(dist_along_path_list[idx]) < 2:
+#             continue
+
+#         dist_along_path = torch.tensor(dist_along_path_list[idx], device=device)
+#         polylines_cur = torch.stack(polylines_list[idx])
+
+#         new_dist_along_path = torch.arange(0, dist_along_path[-1], distance)
+#         new_dist_along_path = torch.cat([new_dist_along_path, dist_along_path[[-1]]])
+
+#         new_x = torch.interp(new_dist_along_path, dist_along_path, polylines_cur[:, 0])
+#         new_y = torch.interp(new_dist_along_path, dist_along_path, polylines_cur[:, 1])
+
+#         new_x_list.append(new_x)
+#         new_y_list.append(new_y)
+
+#         new_polylines = torch.stack((new_x, new_y), dim=-1)
+
+#         polyline_size = int(split_distance / distance)
+#         if new_polylines.shape[0] >= (polyline_size + 1):
+#             padding_size = (new_polylines.shape[0] - (polyline_size + 1)) % polyline_size
+#             final_index = (new_polylines.shape[0] - (polyline_size + 1)) // polyline_size + 1
+#         else:
+#             padding_size = new_polylines.shape[0]
+#             final_index = 0
+
+#         multi_polylines = None
+#         new_heading = torch.atan2(new_polylines[1:, 1] - new_polylines[:-1, 1],
+#                                   new_polylines[1:, 0] - new_polylines[:-1, 0])
+#         new_heading = torch.cat([new_heading, new_heading[-1:]], -1)[..., None]
+#         new_polylines = torch.cat([new_polylines, new_heading], -1)
+
+#         if new_polylines.shape[0] >= (polyline_size + 1):
+#             multi_polylines = new_polylines.unfold(dimension=0, size=polyline_size + 1, step=polyline_size)
+#             multi_polylines = multi_polylines.transpose(1, 2)
+#             multi_polylines = multi_polylines[:, ::5, :]
+
+#         if padding_size >= 3:
+#             last_polyline = new_polylines[final_index * polyline_size:]
+#             last_polyline = last_polyline[torch.linspace(0, last_polyline.shape[0] - 1, steps=3).long()]
+#             if multi_polylines is not None:
+#                 multi_polylines = torch.cat([multi_polylines, last_polyline.unsqueeze(0)], dim=0)
+#             else:
+#                 multi_polylines = last_polyline.unsqueeze(0)
+
+#         if multi_polylines is None:
+#             continue
+#         multi_polylines_list.append(multi_polylines)
+
+#     if len(multi_polylines_list) > 0:
+#         multi_polylines_list = torch.cat(multi_polylines_list, dim=0)
+#     else:
+#         multi_polylines_list = None
+
+#     return multi_polylines_list
+
+
+def average_distance_vectorized(point_set1, centroids):
+    dists = np.sqrt(np.sum((point_set1[:, None, :, :] - centroids[None, :, :, :]) ** 2, axis=-1))
+    return np.mean(dists, axis=2)
+
+
+def assign_clusters(sub_X, centroids):
+    distances = average_distance_vectorized(sub_X, centroids)
+    return np.argmin(distances, axis=1)
+
+
+class TokenProcessor(nn.Module):
+
+    def __init__(self, token_size,
+                 training: bool=False,
+                 predict_motion: bool=False,
+                 predict_state: bool=False,
+                 predict_map: bool=False,
+                 state_token: Dict[str, int]=None, **kwargs):
+        super().__init__()
+
+        module_dir = os.path.dirname(os.path.dirname(__file__))
+        self.agent_token_path = os.path.join(module_dir, f'tokens/agent_vocab_555_s2.pkl')
+        self.map_token_traj_path = os.path.join(module_dir, 'tokens/map_traj_token5.pkl')
+        assert os.path.exists(self.agent_token_path), f"File {self.agent_token_path} not found."
+        assert os.path.exists(self.map_token_traj_path), f"File {self.map_token_traj_path} not found."
+
+        self.training = training
+        self.token_size = token_size
+        self.disable_invalid = not predict_state
+        self.predict_motion = predict_motion
+        self.predict_state = predict_state
+        self.predict_map = predict_map
+
+        # define new special tokens
+        self.bos_token_index = token_size
+        self.eos_token_index = token_size + 1
+        self.invalid_token_index = token_size + 2
+        self.special_token_index = []
+        self._init_token()
+
+        # define agent states
+        self.invalid_state = int(state_token['invalid'])
+        self.valid_state = int(state_token['valid'])
+        self.enter_state = int(state_token['enter'])
+        self.exit_state = int(state_token['exit'])
+
+        self.pl2seed_radius = kwargs.get('pl2seed_radius', None)
+
+        self.noise = False
+        self.disturb = False
+        self.shift = 5
+        self.training = False
+        self.current_step = 10
+
+        # debugging
+        self.debug_data = None
+
+    def forward(self, data):
+        """
+        Each pkl data represents a extracted scenario from raw tfrecord data
+        """
+        data['agent']['av_index'] = data['agent']['av_idx']
+        data = self._tokenize_agent(data)
+        # data = self._tokenize_map(data)
+        del data['city']
+        if 'polygon_is_intersection' in data['map_polygon']:
+            del data['map_polygon']['polygon_is_intersection']
+        if 'route_type' in data['map_polygon']:
+            del data['map_polygon']['route_type']
+
+        av_index = int(data['agent']['av_idx'])
+        data['ego_pos'] = data['agent']['token_pos'][[av_index]]
+        data['ego_heading'] = data['agent']['token_heading'][[av_index]]
+
+        return data
+
+    def _init_token(self):
+
+        agent_token_data = pickle.load(open(self.agent_token_path, 'rb'))
+        for agent_type, token in agent_token_data['token_all'].items():
+            token = torch.tensor(token, dtype=torch.float32)
+            self.register_buffer(f'agent_token_all_{agent_type}', token, persistent=False)  # [n_token, 6, 4, 2]
+
+        map_token_traj = pickle.load(open(self.map_token_traj_path, 'rb'))['traj_src']
+        map_token_traj = torch.tensor(map_token_traj, dtype=torch.float32)
+        self.register_buffer('map_token_traj_src', map_token_traj, persistent=False)  # [n_token, 11 * 2]
+
+        # self.trajectory_token = agent_token_data['token']           # (token_size, 4, 2)
+        # self.trajectory_token_all = agent_token_data['token_all']   # (token_size, shift + 1, 4, 2)
+        # self.map_token = {'traj_src': map_token_traj['traj_src']}
+
+    @staticmethod
+    def clean_heading(valid: torch.Tensor, heading: torch.Tensor) -> torch.Tensor:
+        valid_pairs = valid[:, :-1] & valid[:, 1:]
+        for i in range(heading.shape[1] - 1):
+            heading_diff = torch.abs(wrap_angle(heading[:, i] - heading[:, i + 1]))
+            change_needed = (heading_diff > 1.5) & valid_pairs[:, i]
+            heading[:, i + 1][change_needed] = heading[:, i][change_needed]
+        return heading
+
+    def _extrapolate_agent_to_prev_token_step(self, valid, pos, heading, vel) -> Sequence[torch.Tensor]:
+        # [n_agent], max will give the first True step
+        first_valid_step = torch.max(valid, dim=1).indices
+
+        for i, t in enumerate(first_valid_step):  # extrapolate to previous 5th step.
+            n_step_to_extrapolate = t % self.shift
+            if (t == self.current_step) and (not valid[i, self.current_step - self.shift]):
+                # such that at least one token is valid in the history.
+                n_step_to_extrapolate = self.shift
+
+            if n_step_to_extrapolate > 0:
+                vel[i, t - n_step_to_extrapolate : t] = vel[i, t]
+                valid[i, t - n_step_to_extrapolate : t] = True
+                heading[i, t - n_step_to_extrapolate : t] = heading[i, t]
+
+                for j in range(n_step_to_extrapolate):
+                    pos[i, t - j - 1] = pos[i, t - j] - vel[i, t] * 0.1
+
+        return valid, pos, heading, vel
+
+    def _get_agent_shape(self, agent_type_masks: dict) -> torch.Tensor:
+        agent_shape = 0.
+        for type, type_mask in agent_type_masks.items():
+            if type == 'veh': width = 2.; length = 4.8
+            if type == 'ped': width = 1.; length = 2.
+            if type == 'cyc': width = 1.; length = 1.
+            agent_shape += torch.stack([width * type_mask, length * type_mask], dim=-1)
+
+        return agent_shape
+
+    def _get_token_traj_all(self, agent_type_masks: dict) -> torch.Tensor:
+        token_traj_all = 0.
+        for type, type_mask in agent_type_masks.items():
+            token_traj_all += type_mask[:, None, None, None, None] * (
+                getattr(self, f'agent_token_all_{type}').unsqueeze(0)
+            )
+        return token_traj_all
+
+    def _tokenize_agent(self, data):
+
+        # get raw data
+        valid_mask = data['agent']['valid_mask']                      # [n_agent, n_step]
+        agent_heading = data['agent']['heading']                      # [n_agent, n_step]
+        agent_pos = data['agent']['position'][..., :2].contiguous()   # [n_agent, n_step, 2]
+        agent_vel = data['agent']['velocity']                         # [n_agent, n_step, 2]
+        agent_type = data['agent']['type']
+        agent_category = data['agent']['category']
+
+        n_agent, n_all_step = valid_mask.shape
+
+        agent_type_masks = {
+            "veh": agent_type == 0,
+            "ped": agent_type == 1,
+            "cyc": agent_type == 2,
+        }
+        agent_heading = self.clean_heading(valid_mask, agent_heading)
+        agent_shape = self._get_agent_shape(agent_type_masks)
+        token_traj_all = self._get_token_traj_all(agent_type_masks) 
+        valid_mask, agent_pos, agent_heading, agent_vel = self._extrapolate_agent_to_prev_token_step(
+            valid_mask, agent_pos, agent_heading, agent_vel
+        )
+        token_traj = token_traj_all[:, :, -1, ...]
+        data['agent']['token_traj_all'] = token_traj_all     # [n_agent, n_token, 6, 4, 2]
+        data['agent']['token_traj'] = token_traj             # [n_agent, n_token, 4, 2]
+
+        valid_mask_shift = valid_mask.unfold(1, self.shift + 1, self.shift)
+        token_valid_mask = valid_mask_shift[:, :, 0] * valid_mask_shift[:, :, -1]
+
+        # vehicle_mask = agent_type == 0
+        # cyclist_mask = agent_type == 2
+        # ped_mask = agent_type == 1
+
+        # veh_pos = agent_pos[vehicle_mask, :, :]
+        # veh_valid_mask = valid_mask[vehicle_mask, :]
+        # cyc_pos = agent_pos[cyclist_mask, :, :]
+        # cyc_valid_mask = valid_mask[cyclist_mask, :]
+        # ped_pos = agent_pos[ped_mask, :, :]
+        # ped_valid_mask = valid_mask[ped_mask, :]
+
+        # index: [n_agent, n_step] contour: [n_agent, n_step, 4, 2]
+        token_index, token_contour, token_all = self._match_agent_token(
+            valid_mask, agent_pos, agent_heading, agent_shape, token_traj, None # token_traj_all
+        )
+
+        traj_pos = traj_heading = None
+        if len(token_all) > 0:
+            traj_pos = token_all.mean(dim=3)  # [n_agent, n_step, 6, 2]
+            diff_xy = token_all[..., 0, :] - token_all[..., 3, :]
+            traj_heading = torch.arctan2(diff_xy[..., 1], diff_xy[..., 0])
+        token_pos = token_contour.mean(dim=2)  # [n_agent, n_step, 2]
+        diff_xy = token_contour[:, :, 0, :] - token_contour[:, :, 3, :]
+        token_heading = torch.arctan2(diff_xy[:, :, 1], diff_xy[:, :, 0])
+
+        # token_index: (num_agent, num_timestep // shift) e.g. (49, 18)
+        # token_contour: (num_agent, num_timestep // shift, contour_dim, feat_dim, 2) e.g. (49, 18, 4, 2)
+        # veh_token_index, veh_token_contour = self._match_agent_token(veh_valid_mask, veh_pos, agent_heading[vehicle_mask],
+        #                                                             'veh', agent_shape[vehicle_mask])
+        # ped_token_index, ped_token_contour = self._match_agent_token(ped_valid_mask, ped_pos, agent_heading[ped_mask],
+        #                                                             'ped', agent_shape[ped_mask])
+        # cyc_token_index, cyc_token_contour = self._match_agent_token(cyc_valid_mask, cyc_pos, agent_heading[cyclist_mask],
+        #                                                             'cyc', agent_shape[cyclist_mask])
+
+        # token_index = torch.zeros((agent_pos.shape[0], veh_token_index.shape[1])).to(torch.int64)
+        # token_index[vehicle_mask] = veh_token_index
+        # token_index[ped_mask] = ped_token_index
+        # token_index[cyclist_mask] = cyc_token_index
+
+        # ! compute agent states
+        bos_index = torch.argmax(token_valid_mask.long(), dim=1)
+        eos_index = token_valid_mask.shape[1] - 1 - torch.argmax(torch.flip(token_valid_mask.long(), dims=[1]), dim=1)
+        state_index = torch.ones_like(token_index)  # init with all valid
+        step_index = torch.arange(state_index.shape[1])[None].repeat(state_index.shape[0], 1).to(token_index.device)
+        state_index[step_index == bos_index[:, None]] = self.enter_state
+        state_index[step_index == eos_index[:, None]] = self.exit_state
+        state_index[(step_index < bos_index[:, None]) | (step_index > eos_index[:, None])] = self.invalid_state
+        # ! IMPORTANT: if the last step is exit token, should convert it back to valid token
+        state_index[state_index[:, -1] == self.exit_state, -1] = self.valid_state
+
+        # update token attributions according to state tokens
+        token_valid_mask[state_index == self.enter_state] = False
+        token_pos[state_index == self.invalid_state] = 0.
+        token_heading[state_index == self.invalid_state] = 0.
+        for i in range(self.shift, agent_pos.shape[1], self.shift):
+            is_bos = state_index[:, i // self.shift - 1] == self.enter_state
+            token_pos[is_bos, i // self.shift - 1] = agent_pos[is_bos, i].clone()
+            # token_heading[is_bos, i // self.shift - 1] = agent_heading[is_bos, i].clone()
+        token_index[state_index == self.invalid_state] = -1
+        token_index[state_index == self.enter_state] = -2
+
+        # acc_token_valid_step = torch.concat([torch.zeros_like(token_valid_mask[:, :1]),
+        #                                      torch.cumsum(token_valid_mask.int(), dim=1),
+        #                                      torch.zeros_like(token_valid_mask[:, -1:])], dim=1)
+        # state_index = torch.ones_like(token_index)  # init with all valid
+        # max_valid_index = torch.argmax(acc_token_valid_step, dim=1)
+        # for step in range(1, acc_token_valid_step.shape[1] - 1):
+
+        #     # replace part of motion tokens with special tokens
+        #     is_bos = (acc_token_valid_step[:, step] == 0) & (acc_token_valid_step[:, step + 1] == 1)
+        #     is_eos = (step == max_valid_index) & (step < acc_token_valid_step.shape[1] - 2) & ~is_bos
+        #     is_invalid = ~token_valid_mask[:, step - 1] & ~is_bos & ~is_eos
+
+        #     state_index[is_bos, step - 1] = self.enter_state
+        #     state_index[is_eos, step - 1] = self.exit_state
+        #     state_index[is_invalid, step - 1] = self.invalid_state
+
+        # token_valid_mask[state_index[:, 0] == self.valid_state, 0] = False
+        # state_index[state_index[:, 0] == self.valid_state, 0] = self.enter_state
+
+        # token_contour = torch.zeros((agent_pos.shape[0], veh_token_contour.shape[1],
+        #                              veh_token_contour.shape[2], veh_token_contour.shape[3]))
+        # token_contour[vehicle_mask] = veh_token_contour
+        # token_contour[ped_mask] = ped_token_contour
+        # token_contour[cyclist_mask] = cyc_token_contour
+
+        raw_token_valid_mask = token_valid_mask.clone()
+        if not self.disable_invalid:
+            token_valid_mask = torch.ones_like(token_valid_mask).bool()
+
+        # apply mask
+        # apply_mask = raw_token_valid_mask.sum(dim=-1) > 2
+        # if self.training and os.getenv('AUG_MASK', False):
+        #     aug_mask = torch.randint(0, 2, (raw_token_valid_mask.shape[0],)).to(raw_token_valid_mask).bool()
+        #     apply_mask &= aug_mask
+
+        # remove invalid agents which are outside the range of pl2inva_radius
+        # remove_ina_mask = torch.zeros_like(data['agent']['train_mask'])
+        # if self.pl2seed_radius is not None:
+        #     num_history_token = 1 if self.training else 2 # NOTE: hard code!!!
+        #     av_index = int(data['agent']['av_index'])
+        #     is_invalid = torch.any(state_index[:, :num_history_token] == self.invalid_state, dim=-1)
+        #     ina_bos_mask = (state_index == self.enter_state) & is_invalid[:, None]
+        #     invalid_bos_step = torch.nonzero(ina_bos_mask, as_tuple=False)
+        #     av_bos_pos = token_pos[av_index, invalid_bos_step[:, 1]] # (num_invalid_bos, 2)
+        #     ina_bos_pos = token_pos[invalid_bos_step[:, 0], invalid_bos_step[:, 1]] # (num_invalid_bos, 2)
+        #     distance = torch.sqrt(torch.sum((ina_bos_pos - av_bos_pos) ** 2, dim=-1))
+        #     remove_ina_mask = (distance > self.pl2seed_radius) | (distance < 0.)
+        #     # apply_mask[invalid_bos_step[remove_ina_mask, 0]] = False
+
+        # data['agent']['remove_ina_mask'] = remove_ina_mask
+
+        # apply_mask[int(data['agent']['av_index'])] = True
+        # data['agent']['num_nodes'] = apply_mask.sum()
+
+        # av_id = data['agent']['id'][data['agent']['av_index']]
+        # data['agent']['id'] = [data['agent']['id'][i] for i in range(len(apply_mask)) if apply_mask[i]]
+        # data['agent']['av_index'] = data['agent']['id'].index(av_id)
+        # data['agent']['id'] = torch.tensor(data['agent']['id'], dtype=torch.long)
+
+        # agent_keys = ['valid_mask', 'predict_mask', 'type', 'category', 'position', 'heading', 'velocity', 'shape']
+        # for key in agent_keys:
+        #     if key in data['agent']:
+        #         data['agent'][key] = data['agent'][key][apply_mask]
+
+        # reset agent shapes
+        for i in range(n_agent):
+            bos_shape_index = torch.nonzero(torch.all(data['agent']['shape'][i] != 0., dim=-1))[0]
+            data['agent']['shape'][i, :] = data['agent']['shape'][i, bos_shape_index]
+        if torch.any(torch.all(data['agent']['shape'][i] == 0., dim=-1)):
+            raise ValueError(f"Found invalid shape values.")
+
+        # compute mean height values for each scenario
+        raw_height = data['agent']['position'][:, self.current_step, 2]
+        valid_height = raw_token_valid_mask[:, 1].bool()
+        veh_mean_z = raw_height[agent_type_masks['veh'] & valid_height].mean()
+        ped_mean_z = raw_height[agent_type_masks['ped'] & valid_height].mean().nan_to_num_(veh_mean_z)  # FIXME: hard code
+        cyc_mean_z = raw_height[agent_type_masks['cyc'] & valid_height].mean().nan_to_num_(veh_mean_z)
+
+        # output
+        data['agent']['token_idx'] = token_index
+        data['agent']['state_idx'] = state_index
+        data['agent']['token_contour'] = token_contour
+        data['agent']['traj_pos'] = traj_pos
+        data['agent']['traj_heading'] = traj_heading
+        data['agent']['token_pos'] = token_pos
+        data['agent']['token_heading'] = token_heading
+        data['agent']['agent_valid_mask'] = token_valid_mask # (a, t)
+        data['agent']['raw_agent_valid_mask'] = raw_token_valid_mask
+        data['agent']['raw_height'] = dict(veh=veh_mean_z,
+                                           ped=ped_mean_z,
+                                           cyc=cyc_mean_z)
+        for type in ['veh', 'ped', 'cyc']:
+            data['agent'][f'trajectory_token_{type}'] = getattr(
+                self, f'agent_token_all_{type}')  # [n_token, 6, 4, 2]
+
+        return data
+
+    def _match_agent_token(self, valid_mask, pos, heading, shape, token_traj, token_traj_all=None):
+        """
+        Parameters:
+        valid_mask (torch.Tensor): Validity mask for agents over time. Shape: (n_agent, n_step)
+        pos (torch.Tensor): Positions of agents at each time step. Shape: (n_agent, n_step, 3)
+        heading (torch.Tensor): Headings of agents at each time step. Shape: (n_agent, n_step)
+        shape (torch.Tensor): Shape information of agents. Shape: (n_agent, 3)
+        token_traj (torch.Tensor): Token trajectories for agents. Shape: (n_agent, n_token, 4, 2)
+        token_traj_all (torch.Tensor): Token trajectories for all agents. Shape: (n_agnet, n_token_all, n_contour, 4, 2)
+
+        Returns:
+        tuple: Contains token indices and contours for agents.
+        """
+
+        n_agent, n_step = valid_mask.shape
+
+        # agent_token_src = self.trajectory_token[category]
+        # if self.shift <= 2:
+        #     if category == 'veh':
+        #         width = 1.0
+        #         length = 2.4
+        #     elif category == 'cyc':
+        #         width = 0.5
+        #         length = 1.5
+        #     else:
+        #         width = 0.5
+        #         length = 0.5
+        # else:
+        #     if category == 'veh':
+        #         width = 2.0
+        #         length = 4.8
+        #     elif category == 'cyc':
+        #         width = 1.0
+        #         length = 2.0
+        #     else:
+        #         width = 1.0
+        #         length = 1.0
+
+        _, n_token, token_contour_dim, feat_dim = token_traj.shape
+        # agent_token_src = agent_token_src.reshape(1, token_num * token_contour_dim, feat_dim).repeat(agent_num, 0)
+
+        token_index_list = []
+        token_contour_list = []
+        token_all = []
+
+        prev_heading = heading[:, 0]
+        prev_pos = pos[:, 0]
+        prev_token_idx = None
+        for i in range(self.shift, n_step, self.shift):  # [5, 10, 15, ..., 90]
+            _valid_mask = valid_mask[:, i - self.shift] & valid_mask[:, i]
+            _invalid_mask = ~_valid_mask
+
+            # transformation
+            theta = prev_heading
+            cos, sin = theta.cos(), theta.sin()
+            rot_mat = theta.new_zeros(n_agent, 2, 2)
+            rot_mat[:, 0, 0] = cos
+            rot_mat[:, 0, 1] = sin
+            rot_mat[:, 1, 0] = -sin
+            rot_mat[:, 1, 1] = cos
+            agent_token_world = torch.bmm(token_traj.flatten(1, 2), rot_mat).reshape(*token_traj.shape)
+            agent_token_world += prev_pos[:, None, None, :]
+
+            cur_contour = cal_polygon_contour(pos[:, i], heading[:, i], shape)  # [n_agent, 4, 2]
+            agent_token_index = torch.argmin(
+                torch.norm(agent_token_world - cur_contour[:, None, ...], dim=-1).sum(-1), dim=-1
+            )
+            agent_token_contour = agent_token_world[torch.arange(n_agent), agent_token_index]  # [n_agent, 4, 2]
+            # agent_token_index = torch.from_numpy(np.argmin(
+            #     np.mean(np.sqrt(np.sum((cur_contour[:, None, ...] - agent_token_world.numpy()) ** 2, axis=-1)), axis=2),
+            #     axis=-1))
+
+            # except for the first timestep TODO
+            if prev_token_idx is not None and self.noise:
+                same_idx = prev_token_idx == agent_token_index
+                same_idx[:] = True
+                topk_indices = np.argsort(
+                    np.mean(np.sqrt(np.sum((cur_contour[:, None, ...] - agent_token_world.numpy()) ** 2, axis=-1)),
+                            axis=2), axis=-1)[:, :5]
+                sample_topk = np.random.choice(range(0, topk_indices.shape[1]), topk_indices.shape[0])
+                agent_token_index[same_idx] = \
+                    torch.from_numpy(topk_indices[np.arange(topk_indices.shape[0]), sample_topk])[same_idx]
+
+            # update prev_heading
+            prev_heading = heading[:, i].clone()
+            diff_xy = agent_token_contour[:, 0] - agent_token_contour[:, 3]
+            prev_heading[_valid_mask] = torch.arctan2(diff_xy[:, 1], diff_xy[:, 0])[_valid_mask]
+
+            # update prev_pos
+            prev_pos = pos[:, i].clone()
+            prev_pos[_valid_mask] = agent_token_contour.mean(dim=1)[_valid_mask]
+
+            prev_token_idx = agent_token_index
+            token_index_list.append(agent_token_index)
+            token_contour_list.append(agent_token_contour)
+
+            # calculate tokenized trajectory
+            if token_traj_all is not None:
+                agent_token_all_world = torch.bmm(token_traj_all.flatten(1, 3), rot_mat).reshape(*token_traj_all.shape)
+                agent_token_all_world += prev_pos[:, None, None, None, :]
+                agent_token_all = agent_token_all_world[torch.arange(n_agent), agent_token_index]  # [n_agent, 6, 4, 2]
+                token_all.append(agent_token_all)
+
+        token_index = torch.stack(token_index_list, dim=1)      # [n_agent, n_step]
+        token_contour = torch.stack(token_contour_list, dim=1)  # [n_agent, n_step, 4, 2]
+        if len(token_all) > 0:
+            token_all = torch.stack(token_all, dim=1)  # [n_agent, n_step, 6, 4, 2]
+
+        # sanity check
+        assert tuple(token_index.shape) == (n_agent, n_step // self.shift), \
+            f'Invalid token_index shape, got {token_index.shape}'
+        assert tuple(token_contour.shape )== (n_agent, n_step // self.shift, token_contour_dim, feat_dim), \
+            f'Invalid token_contour shape, got {token_contour.shape}'
+
+        # extra matching
+        # if not self.training:
+        #     theta = heading[extra_mask, self.current_step - 1]
+        #     prev_pos = pos[extra_mask, self.current_step - 1]
+        #     cur_pos = pos[extra_mask, self.current_step]
+        #     cur_heading = heading[extra_mask, self.current_step]
+        #     cos, sin = theta.cos(), theta.sin()
+        #     rot_mat = theta.new_zeros(extra_mask.sum(), 2, 2)
+        #     rot_mat[:, 0, 0] = cos
+        #     rot_mat[:, 0, 1] = sin
+        #     rot_mat[:, 1, 0] = -sin
+        #     rot_mat[:, 1, 1] = cos
+        #     agent_token_world = torch.bmm(torch.from_numpy(token_last).to(torch.float), rot_mat).reshape(
+        #         extra_mask.sum(), token_num, token_contour_dim, feat_dim)
+        #     agent_token_world += prev_pos[:, None, None, :]
+
+        #     cur_contour = cal_polygon_contour(cur_pos[:, 0], cur_pos[:, 1], cur_heading, width, length)
+        #     agent_token_index = torch.from_numpy(np.argmin(
+        #         np.mean(np.sqrt(np.sum((cur_contour[:, None, ...] - agent_token_world.numpy()) ** 2, axis=-1)), axis=2),
+        #         axis=-1))
+        #     token_contour_select = agent_token_world[torch.arange(extra_mask.sum()), agent_token_index]
+
+        #     token_index[extra_mask, 1] = agent_token_index
+        #     token_contour[extra_mask, 1] = token_contour_select
+
+        return token_index, token_contour, token_all
+
+    @staticmethod
+    def _tokenize_map(data):
+
+        data['map_polygon']['type'] = data['map_polygon']['type'].to(torch.uint8)
+        data['map_point']['type'] = data['map_point']['type'].to(torch.uint8)
+        pt2pl = data[('map_point', 'to', 'map_polygon')]['edge_index']
+        pt_type = data['map_point']['type'].to(torch.uint8)
+        pt_side = torch.zeros_like(pt_type)
+        pt_pos = data['map_point']['position'][:, :2]
+        data['map_point']['orientation'] = wrap_angle(data['map_point']['orientation'])
+        pt_heading = data['map_point']['orientation']
+        split_polyline_type = []
+        split_polyline_pos = []
+        split_polyline_theta = []
+        split_polyline_side = []
+        pl_idx_list = []
+        split_polygon_type = []
+        data['map_point']['type'].unique()
+
+        for i in sorted(np.unique(pt2pl[1])):  # number of polygons in the scenario
+            index = pt2pl[0, pt2pl[1] == i]  # index of points which belongs to i-th polygon
+            polygon_type = data['map_polygon']["type"][i]
+            cur_side = pt_side[index]
+            cur_type = pt_type[index]
+            cur_pos = pt_pos[index]
+            cur_heading = pt_heading[index]
+
+            for side_val in np.unique(cur_side):
+                for type_val in np.unique(cur_type):
+                    if type_val == 13:
+                        continue
+                    indices = np.where((cur_side == side_val) & (cur_type == type_val))[0]
+                    if len(indices) <= 2:
+                        continue
+                    split_polyline = interplating_polyline(cur_pos[indices].numpy(), cur_heading[indices].numpy())
+                    if split_polyline is None:
+                        continue
+                    new_cur_type = cur_type[indices][0]
+                    new_cur_side = cur_side[indices][0]
+                    map_polygon_type = polygon_type.repeat(split_polyline.shape[0])
+                    new_cur_type = new_cur_type.repeat(split_polyline.shape[0])
+                    new_cur_side = new_cur_side.repeat(split_polyline.shape[0])
+                    cur_pl_idx = torch.Tensor([i])
+                    new_cur_pl_idx = cur_pl_idx.repeat(split_polyline.shape[0])
+                    split_polyline_pos.append(split_polyline[..., :2])
+                    split_polyline_theta.append(split_polyline[..., 2])
+                    split_polyline_type.append(new_cur_type)
+                    split_polyline_side.append(new_cur_side)
+                    pl_idx_list.append(new_cur_pl_idx)
+                    split_polygon_type.append(map_polygon_type)
+
+        split_polyline_pos = torch.cat(split_polyline_pos, dim=0)
+        split_polyline_theta = torch.cat(split_polyline_theta, dim=0)
+        split_polyline_type = torch.cat(split_polyline_type, dim=0)
+        split_polyline_side = torch.cat(split_polyline_side, dim=0)
+        split_polygon_type = torch.cat(split_polygon_type, dim=0)
+        pl_idx_list = torch.cat(pl_idx_list, dim=0)
+
+        data['map_save'] = {}
+        data['pt_token'] = {}
+        data['map_save']['traj_pos'] = split_polyline_pos
+        data['map_save']['traj_theta'] = split_polyline_theta[:, 0]  # torch.arctan2(vec[:, 1], vec[:, 0])
+        data['map_save']['pl_idx_list'] = pl_idx_list
+        data['pt_token']['type'] = split_polyline_type
+        data['pt_token']['side'] = split_polyline_side
+        data['pt_token']['pl_type'] = split_polygon_type
+        data['pt_token']['num_nodes'] = split_polyline_pos.shape[0]
+
+        return data

seq_10_150_3_3_encode_occ_separate_offsets_bs8_128_no_seqindex_long_lastvalid_ablation_grid/backups/dev/datasets/scalable_dataset.py

@@ -0,0 +1,276 @@
+import os
+import pickle
+import torch
+import json
+import pytorch_lightning as pl
+import pandas as pd
+from tqdm import tqdm
+from torch_geometric.data import HeteroData, Dataset
+from torch_geometric.transforms import BaseTransform
+from torch_geometric.loader import DataLoader
+from typing import Callable, Dict, List, Optional
+
+from dev.datasets.preprocess import TokenProcessor
+
+
+class MultiDataset(Dataset):
+    def __init__(self,
+                 split: str,
+                 raw_dir: List[str] = None,
+                 transform: Optional[Callable] = None,
+                 tfrecord_dir: Optional[str] = None,
+                 token_size=512,
+                 predict_motion: bool=False,
+                 predict_state: bool=False,
+                 predict_map: bool=False,
+                #  state_token: Dict[str, int]=None,
+                #  pl2seed_radius: float=None,
+                 buffer_size: int=128,
+                 logger=None) -> None:
+
+        self.disable_invalid = not predict_state
+        self.predict_motion = predict_motion
+        self.predict_state = predict_state
+        self.predict_map = predict_map
+        self.logger = logger
+        if split not in ('train', 'val', 'test'):
+            raise ValueError(f'{split} is not a valid split')
+        self.training = split == 'train'
+        self.buffer_size = buffer_size
+        self._tfrecord_dir = tfrecord_dir
+        self.logger.debug('Starting loading dataset')
+
+        raw_dir = os.path.expanduser(os.path.normpath(raw_dir))
+        self._raw_files = sorted(os.listdir(raw_dir))
+
+        # for debugging
+        if int(os.getenv('OVERFIT', 0)):
+            # if self.training:
+            #     # scenario_id = ['74ad7b76d5906d39', '13596229fd8cdb7e', '1d73db1fc42be3bf', '1351ea8b8333ddcb']
+            #     self._raw_files = ['74ad7b76d5906d39.pkl'] + self._raw_files[:9]
+            # else:
+            #     self._raw_files = self._raw_files[:10]
+            self._raw_files = self._raw_files[:1]
+            # self._raw_files = ['1002fdc9826fc6d1.pkl']
+
+        # load meta infos and do filter
+        json_path = '/u/xiuyu/work/dev4/data/waymo_processed/meta_infos.json'
+        label = 'training' if split == 'train' else ('validation' if split == 'val' else split)
+        self.meta_infos = json.load(open(json_path, 'r', encoding='utf-8'))[label]
+        self.logger.debug(f"Loaded meta infos from {json_path}")
+        self.available_scenarios = list(self.meta_infos.keys())
+        # self._raw_files = list(tqdm(filter(lambda fn: (
+        #     scenario_id := fn.removesuffix('.pkl') in self.available_scenarios and
+        #     8 <= self.meta_infos[scenario_id]['num_agents'] < self.buffer_size
+        # ), self._raw_files), leave=False))
+        df = pd.DataFrame.from_dict(self.meta_infos, orient='index')
+        available_scenarios_set = set(self.available_scenarios)
+        df_filtered = df[(df.index.isin(available_scenarios_set)) & (df['num_agents'] >= 8) & (df['num_agents'] < self.buffer_size)]
+        valid_scenarios = set(df_filtered.index)
+        self._raw_files = list(tqdm(filter(lambda fn: fn.removesuffix('.pkl') in valid_scenarios, self._raw_files), leave=False))
+        if len(self._raw_files) <= 0:
+            raise RuntimeError(f'Invalid number of data {len(self._raw_files)}!')
+        self._raw_paths = list(map(lambda fn: os.path.join(raw_dir, fn), self._raw_files))
+
+        self.logger.debug(f"The number of {split} dataset is {len(self._raw_paths)}")
+        self.logger.debug(f"The buffer size is {self.buffer_size}")
+        # self.token_processor = TokenProcessor(token_size,
+        #                                       training=self.training,
+        #                                       predict_motion=self.predict_motion,
+        #                                       predict_state=self.predict_state,
+        #                                       predict_map=self.predict_map,
+        #                                       state_token=state_token,
+        #                                       pl2seed_radius=pl2seed_radius) # 2048
+        self.logger.debug(f"The used token size is {token_size}.")
+        super().__init__(transform=transform, pre_transform=None, pre_filter=None)
+
+    def len(self) -> int:
+        return len(self._raw_paths)
+
+    def get(self, idx: int):
+        """
+        Load pkl file (each represents a 91s scenario for waymo dataset)
+        """
+        with open(self._raw_paths[idx], 'rb') as handle:
+            data = pickle.load(handle)
+
+        if self._tfrecord_dir is not None:
+            data['tfrecord_path'] = os.path.join(self._tfrecord_dir, f"{data['scenario_id']}.tfrecords")
+
+        # data = self.token_processor.preprocess(data)
+        return data
+
+
+class WaymoTargetBuilder(BaseTransform):
+
+    def __init__(self,
+                 num_historical_steps: int,
+                 num_future_steps: int,
+                 max_num: int,
+                 training: bool=False) -> None:
+
+        self.max_num = max_num
+        self.num_historical_steps = num_historical_steps
+        self.num_future_steps = num_future_steps
+        self.step_current = num_historical_steps - 1
+        self.training = training
+
+    def _score_trained_agents(self, data):
+        pos = data['agent']['position']
+        av_index = torch.where(data['agent']['role'][:, 0])[0].item()
+        distance = torch.norm(pos - pos[av_index], dim=-1)
+
+        # we do not believe the perception out of range of 150 meters
+        data['agent']['valid_mask'] &= distance < 150
+
+        # we do not predict vehicle too far away from ego car
+        role_train_mask = data['agent']['role'].any(-1)
+        extra_train_mask = (distance[:, self.step_current] < 100) & (
+            data['agent']['valid_mask'][:, self.step_current + 1 :].sum(-1) >= 5
+        )
+
+        train_mask = extra_train_mask | role_train_mask
+        if train_mask.sum() > self.max_num:  # too many vehicle
+            _indices = torch.where(extra_train_mask & ~role_train_mask)[0]
+            selected_indices = _indices[
+                torch.randperm(_indices.size(0))[: self.max_num - role_train_mask.sum()]
+            ]
+            data['agent']['train_mask'] = role_train_mask
+            data['agent']['train_mask'][selected_indices] = True
+        else:
+            data['agent']['train_mask'] = train_mask  # [n_agent]
+
+        return data
+
+    def __call__(self, data) -> HeteroData:
+
+        if self.training:
+            self._score_trained_agents(data)
+
+        data = TokenProcessor._tokenize_map(data)
+        # data keys: dict_keys(['scenario_id', 'agent', 'map_polygon', 'map_point', ('map_point', 'to', 'map_polygon'), ('map_polygon', 'to', 'map_polygon'), 'map_save', 'pt_token'])
+        return HeteroData(data)
+
+
+class MultiDataModule(pl.LightningDataModule):
+    transforms = {
+        'WaymoTargetBuilder': WaymoTargetBuilder,
+    }
+
+    dataset = {
+        'scalable': MultiDataset,
+    }
+
+    def __init__(self,
+                 root: str,
+                 train_batch_size: int,
+                 val_batch_size: int,
+                 test_batch_size: int,
+                 shuffle: bool = False,
+                 num_workers: int = 0,
+                 pin_memory: bool = True,
+                 persistent_workers: bool = True,
+                 train_raw_dir: Optional[str] = None,
+                 val_raw_dir: Optional[str] = None,
+                 test_raw_dir: Optional[str] = None,
+                 train_processed_dir: Optional[str] = None,
+                 val_processed_dir: Optional[str] = None,
+                 test_processed_dir: Optional[str] = None,
+                 val_tfrecords_splitted: Optional[str] = None,
+                 transform: Optional[str] = None,
+                 dataset: Optional[str] = None,
+                 num_historical_steps: int = 50,
+                 num_future_steps: int = 60,
+                 processor='ntp',
+                 token_size=512,
+                 predict_motion: bool=False,
+                 predict_state: bool=False,
+                 predict_map: bool=False,
+                 state_token: Dict[str, int]=None,
+                 pl2seed_radius: float=None,
+                 max_num: int=32,
+                 buffer_size: int=256,
+                 logger=None,
+                 **kwargs) -> None:
+
+        super(MultiDataModule, self).__init__()
+        self.root = root
+        self.dataset_class = dataset
+        self.train_batch_size = train_batch_size
+        self.val_batch_size = val_batch_size
+        self.test_batch_size = test_batch_size
+        self.shuffle = shuffle
+        self.num_workers = num_workers
+        self.pin_memory = pin_memory
+        self.persistent_workers = persistent_workers and num_workers > 0
+        self.train_raw_dir = train_raw_dir
+        self.val_raw_dir = val_raw_dir
+        self.test_raw_dir = test_raw_dir
+        self.train_processed_dir = train_processed_dir
+        self.val_processed_dir = val_processed_dir
+        self.test_processed_dir = test_processed_dir
+        self.val_tfrecords_splitted = val_tfrecords_splitted
+        self.processor = processor
+        self.token_size = token_size
+        self.predict_motion = predict_motion
+        self.predict_state = predict_state
+        self.predict_map = predict_map
+        self.state_token = state_token
+        self.pl2seed_radius = pl2seed_radius
+        self.buffer_size = buffer_size
+        self.logger = logger
+
+        self.train_transform = MultiDataModule.transforms[transform](num_historical_steps,
+                                                                     num_future_steps,
+                                                                     max_num=max_num,
+                                                                     training=True)
+        self.val_transform = MultiDataModule.transforms[transform](num_historical_steps,
+                                                                   num_future_steps,
+                                                                   max_num=max_num,
+                                                                   training=False)
+
+    def setup(self, stage: Optional[str] = None) -> None:
+        general_params = dict(token_size=self.token_size,
+                              predict_motion=self.predict_motion,
+                              predict_state=self.predict_state,
+                              predict_map=self.predict_map,
+                              buffer_size=self.buffer_size,
+                              logger=self.logger)
+
+        if stage == 'fit' or stage is None:
+            self.train_dataset = MultiDataModule.dataset[self.dataset_class](split='train',
+                                                                             raw_dir=self.train_raw_dir,
+                                                                             transform=self.train_transform,
+                                                                             **general_params)
+            self.val_dataset = MultiDataModule.dataset[self.dataset_class](split='val',
+                                                                           raw_dir=self.val_raw_dir,
+                                                                           transform=self.val_transform,
+                                                                           tfrecord_dir=self.val_tfrecords_splitted,
+                                                                           **general_params)
+        if stage == 'validate':
+            self.val_dataset = MultiDataModule.dataset[self.dataset_class](split='val',
+                                                                           raw_dir=self.val_raw_dir,
+                                                                           transform=self.val_transform,
+                                                                           tfrecord_dir=self.val_tfrecords_splitted,
+                                                                           **general_params)
+        if stage == 'test':
+            self.test_dataset = MultiDataModule.dataset[self.dataset_class](split='test',
+                                                                            raw_dir=self.test_raw_dir,
+                                                                            transform=self.val_transform,
+                                                                            tfrecord_dir=self.val_tfrecords_splitted,
+                                                                            **general_params)
+
+    def train_dataloader(self):
+        return DataLoader(self.train_dataset, batch_size=self.train_batch_size, shuffle=self.shuffle,
+                          num_workers=self.num_workers, pin_memory=self.pin_memory,
+                          persistent_workers=self.persistent_workers)
+
+    def val_dataloader(self):
+        return DataLoader(self.val_dataset, batch_size=self.val_batch_size, shuffle=False,
+                          num_workers=self.num_workers, pin_memory=self.pin_memory,
+                          persistent_workers=self.persistent_workers)
+
+    def test_dataloader(self):
+        return DataLoader(self.test_dataset, batch_size=self.test_batch_size, shuffle=False,
+                          num_workers=self.num_workers, pin_memory=self.pin_memory,
+                          persistent_workers=self.persistent_workers)

seq_10_150_3_3_encode_occ_separate_offsets_bs8_128_no_seqindex_long_lastvalid_ablation_grid/backups/dev/metrics/box_utils.py

@@ -0,0 +1,113 @@
+import torch
+from torch import Tensor
+
+
+def get_yaw_rotation_2d(yaw):
+    """
+    Gets a 2D rotation matrix given a yaw angle.
+
+    Args:
+        yaw: torch.Tensor, rotation angle in radians. Can be any shape except empty.
+
+    Returns:
+        rotation: torch.Tensor with shape [..., 2, 2], where `...` matches input shape.
+    """
+    cos_yaw = torch.cos(yaw)
+    sin_yaw = torch.sin(yaw)
+
+    rotation = torch.stack([
+        torch.stack([cos_yaw, -sin_yaw], dim=-1),
+        torch.stack([sin_yaw,  cos_yaw], dim=-1),
+    ], dim=-2)  # Shape: [..., 2, 2]
+
+    return rotation
+
+
+def get_yaw_rotation(yaw):
+    """
+    Computes a 3D rotation matrix given a yaw angle (rotation around the Z-axis).
+
+    Args:
+        yaw: torch.Tensor of any shape, representing yaw angles in radians.
+
+    Returns:
+        rotation: torch.Tensor of shape [input_shape, 3, 3], representing the rotation matrices.
+    """
+    cos_yaw = torch.cos(yaw)
+    sin_yaw = torch.sin(yaw)
+    ones = torch.ones_like(yaw)
+    zeros = torch.zeros_like(yaw)
+
+    return torch.stack([
+        torch.stack([cos_yaw, -sin_yaw, zeros], dim=-1),
+        torch.stack([sin_yaw, cos_yaw, zeros], dim=-1),
+        torch.stack([zeros, zeros, ones], dim=-1),
+    ], dim=-2)
+
+
+def get_transform(rotation, translation):
+    """
+    Combines an NxN rotation matrix and an Nx1 translation vector into an (N+1)x(N+1) transformation matrix.
+
+    Args:
+        rotation: torch.Tensor of shape [..., N, N], representing rotation matrices.
+        translation: torch.Tensor of shape [..., N], representing translation vectors.
+                    This must have the same dtype as rotation.
+
+    Returns:
+        transform: torch.Tensor of shape [..., (N+1), (N+1)], representing the transformation matrices.
+                   This has the same dtype as rotation.
+    """
+    # [..., N, 1]
+    translation_n_1 = translation.unsqueeze(-1)
+    
+    # [..., N, N+1] - Combine rotation and translation
+    transform = torch.cat([rotation, translation_n_1], dim=-1)
+    
+    # [..., N] - Create the last row, which is [0, 0, ..., 0, 1]
+    last_row = torch.zeros_like(translation)
+    last_row = torch.cat([last_row, torch.ones_like(last_row[..., :1])], dim=-1)
+    
+    # [..., N+1, N+1] - Append the last row to form the final transformation matrix
+    transform = torch.cat([transform, last_row.unsqueeze(-2)], dim=-2)
+    
+    return transform
+
+
+def get_upright_3d_box_corners(boxes: Tensor):
+    """
+    Given a set of upright 3D bounding boxes, return its 8 corner points.
+
+    Args:
+        boxes: torch.Tensor [N, 7]. The inner dims are [center{x,y,z}, length, width,
+               height, heading].
+
+    Returns:
+        corners: torch.Tensor [N, 8, 3].
+    """
+    center_x, center_y, center_z, length, width, height, heading = boxes.unbind(dim=-1)
+
+    # Compute rotation matrix [N, 3, 3]
+    rotation = get_yaw_rotation(heading)
+
+    # Translation [N, 3]
+    translation = torch.stack([center_x, center_y, center_z], dim=-1)
+
+    l2, w2, h2 = length * 0.5, width * 0.5, height * 0.5
+
+    # Define the 8 corners in local coordinates [N, 8, 3]
+    corners_local = torch.stack([
+        torch.stack([ l2,  w2, -h2], dim=-1),
+        torch.stack([-l2,  w2, -h2], dim=-1),
+        torch.stack([-l2, -w2, -h2], dim=-1),
+        torch.stack([ l2, -w2, -h2], dim=-1),
+        torch.stack([ l2,  w2,  h2], dim=-1),
+        torch.stack([-l2,  w2,  h2], dim=-1),
+        torch.stack([-l2, -w2,  h2], dim=-1),
+        torch.stack([ l2, -w2,  h2], dim=-1),
+    ], dim=1)  # Shape: [N, 8, 3]
+
+    # Rotate and translate the corners
+    corners = torch.einsum('n i j, n k j -> n k i', rotation, corners_local) + translation.unsqueeze(1)
+
+    return corners

seq_10_150_3_3_encode_occ_separate_offsets_bs8_128_no_seqindex_long_lastvalid_ablation_grid/backups/dev/metrics/compute_metrics.py

@@ -0,0 +1,1812 @@
+# ! Metrics Calculation
+import concurrent.futures
+import os
+import torch
+import tensorflow as tf
+import collections
+import dataclasses
+import fnmatch
+import json
+import pandas as pd
+import pickle
+import copy
+import concurrent
+import multiprocessing
+from torch_geometric.utils import degree
+from functools import partial
+from dataclasses import dataclass, field
+from tqdm import tqdm
+from argparse import ArgumentParser
+from torch import Tensor
+from google.protobuf import text_format
+from torchmetrics import Metric
+from typing import Optional, Sequence, List, Dict
+
+from waymo_open_dataset.utils.sim_agents import submission_specs
+
+from dev.utils.visualization import safe_run
+from dev.utils.func import CONSOLE
+from dev.datasets.scalable_dataset import WaymoTargetBuilder
+from dev.datasets.preprocess import TokenProcessor, SHIFT, AGENT_STATE
+from dev.metrics import trajectory_features, interact_features, map_features, placement_features
+from dev.metrics.protos import scenario_pb2, long_metrics_pb2
+
+
+_METRIC_FIELD_NAMES_BY_BUCKET = {
+    'kinematic': [
+        'linear_speed', 'linear_acceleration',
+        'angular_speed', 'angular_acceleration',
+    ],
+    'interactive': [
+        'distance_to_nearest_object', 'collision_indication',
+        'time_to_collision',
+    ],
+    'map_based': [
+        # 'distance_to_road_edge', 'offroad_indication'
+    ],
+    'placement_based': [
+        'num_placement', 'num_removement',
+        'distance_placement', 'distance_removement',
+    ]
+}
+_METRIC_FIELD_NAMES = (
+    _METRIC_FIELD_NAMES_BY_BUCKET['kinematic'] +
+    _METRIC_FIELD_NAMES_BY_BUCKET['interactive'] +
+    _METRIC_FIELD_NAMES_BY_BUCKET['map_based'] +
+    _METRIC_FIELD_NAMES_BY_BUCKET['placement_based']
+)
+
+
+""" Help Functions """
+
+def _arg_gather(tensor: Tensor, reference_tensor: Tensor) -> Tensor:
+    """Finds corresponding indices in `tensor` for each element in `reference_tensor`.
+
+    Args:
+        tensor: A 1D tensor without repetitions.
+        reference_tensor: A 1D tensor containing items from `tensor`.
+
+    Returns:
+        A tensor of indices such that `tensor[indices] == reference_tensor`.
+    """
+    assert tensor.ndim == 1, "tensor must be 1D"
+    assert reference_tensor.ndim == 1, "reference_tensor must be 1D"
+    
+    # Create the comparison matrix
+    bit_mask = tensor[None, :] == reference_tensor[:, None]  # Shape: [len(reference_tensor), len(tensor)]
+    
+    # Count the matches along `tensor` dimension
+    bit_mask_sum = bit_mask.int().sum(dim=1)
+    
+    if (bit_mask_sum < 1).any():
+        raise ValueError(
+            'Some items in `reference_tensor` are missing from `tensor`: '
+            f'\n{reference_tensor} \nvs. \n{tensor}.'
+        )
+    
+    if (bit_mask_sum > 1).any():
+        raise ValueError('Some items in `tensor` are repeated.')
+    
+    # Compute indices
+    indices = torch.matmul(bit_mask.int(), torch.arange(tensor.shape[0], dtype=torch.int32))
+    return indices
+
+
+def is_valid_sim_agent(track: scenario_pb2.Track) -> bool: # type: ignore
+    """Checks if the object needs to be resimulated as a sim agent.
+
+    For the Sim Agents challenge, every object that is valid at the
+    `current_time_index` step (here hardcoded to 10) needs to be resimulated.
+
+    Args:
+    track: A track proto for a single object.
+
+    Returns:
+    A boolean flag, True if the object needs to be resimulated, False otherwise.
+    """
+    return track.states[submission_specs.CURRENT_TIME_INDEX].valid
+
+
+def get_sim_agent_ids(
+    scenario: scenario_pb2.Scenario) -> Sequence[int]: # type: ignore
+    """Returns the list of object IDs that needs to be resimulated.
+
+    Internally calls `is_valid_sim_agent` to verify the simulation criteria,
+    i.e. is the object valid at `current_time_index`.
+
+    Args:
+    scenario: The Scenario proto containing the data.
+
+    Returns:
+    A list of int IDs, containing all the objects that need to be simulated.
+    """
+    object_ids = []
+    for track in scenario.tracks:
+        if is_valid_sim_agent(track):
+            object_ids.append(track.id)
+    return object_ids
+
+
+def get_evaluation_agent_ids(
+    scenario: scenario_pb2.Scenario) -> Sequence[int]: # type: ignore
+    # Start with the AV object.
+    object_ids = {scenario.tracks[scenario.sdc_track_index].id}
+    # Add the `tracks_to_predict` objects.
+    for required_prediction in scenario.tracks_to_predict:
+        object_ids.add(scenario.tracks[required_prediction.track_index].id)
+    return sorted(object_ids)
+
+
+""" Base Data Classes s"""
+
+@dataclass(frozen=True)
+class ObjectTrajectories:
+
+    x: Tensor
+    y: Tensor
+    z: Tensor
+    heading: Tensor
+    length: Tensor
+    width: Tensor
+    height: Tensor
+    valid: Tensor
+    object_id: Tensor
+    object_type: Tensor
+
+    state: Optional[Tensor] = None
+    token_pos: Optional[Tensor] = None
+    token_heading: Optional[Tensor] = None
+    token_valid: Optional[Tensor] = None
+    processed_object_id: Optional[Tensor] = None
+    av_id: Optional[int] = None
+    processed_av_id: Optional[int] = None
+
+    def slice_time(self, start_index: int = 0, end_index: Optional[int] = None):
+        return ObjectTrajectories(
+            x=self.x[..., start_index:end_index],
+            y=self.y[..., start_index:end_index],
+            z=self.z[..., start_index:end_index],
+            heading=self.heading[..., start_index:end_index],
+            length=self.length[..., start_index:end_index],
+            width=self.width[..., start_index:end_index],
+            height=self.height[..., start_index:end_index],
+            valid=self.valid[..., start_index:end_index],
+            object_id=self.object_id,
+            object_type=self.object_type,
+
+            # these properties can only come from processed file
+            state=self.state,
+            token_pos=self.token_pos,
+            token_heading=self.token_heading,
+            token_valid=self.token_valid,
+            processed_object_id=self.processed_object_id,
+            av_id=self.av_id,
+            processed_av_id=self.processed_av_id,
+        )
+
+    def gather_objects(self, object_indices: Tensor):
+        assert object_indices.ndim == 1, "object_indices must be 1D"
+        return ObjectTrajectories(
+            x=torch.index_select(self.x, dim=-2, index=object_indices),
+            y=torch.index_select(self.y, dim=-2, index=object_indices),
+            z=torch.index_select(self.z, dim=-2, index=object_indices),
+            heading=torch.index_select(self.heading, dim=-2, index=object_indices),
+            length=torch.index_select(self.length, dim=-2, index=object_indices),
+            width=torch.index_select(self.width, dim=-2, index=object_indices),
+            height=torch.index_select(self.height, dim=-2, index=object_indices),
+            valid=torch.index_select(self.valid, dim=-2, index=object_indices),
+            object_id=torch.index_select(self.object_id, dim=-1, index=object_indices),
+            object_type=torch.index_select(self.object_type, dim=-1, index=object_indices),
+
+            # these properties can only come from processed file
+            state=self.state,
+            token_pos=self.token_pos,
+            token_heading=self.token_heading,
+            token_valid=self.token_valid,
+            processed_object_id=self.processed_object_id,
+            av_id=self.av_id,
+            processed_av_id=self.processed_av_id,
+        )
+
+    def gather_objects_by_id(self, object_ids: tf.Tensor):
+        indices = _arg_gather(self.object_id, object_ids)
+        return self.gather_objects(indices)
+
+    @classmethod
+    def _get_init_dict_from_processed(cls, scenario: dict):
+        """Load from processed pkl data"""
+        position = scenario['agent']['position']
+        heading = scenario['agent']['heading']
+        shape = scenario['agent']['shape']
+        object_ids = scenario['agent']['id']
+        object_types = scenario['agent']['type']
+        valid = scenario['agent']['valid_mask']
+
+        init_dict = dict(x=position[..., 0],
+                         y=position[..., 1],
+                         z=position[..., 2],
+                         heading=heading,
+                         length=shape[..., 0],
+                         width=shape[..., 1],
+                         height=shape[..., 2],
+                         valid=valid,
+                         object_ids=object_ids,
+                         object_types=object_types)
+
+        return init_dict
+
+    @classmethod
+    def _get_init_dict_from_raw(cls,
+                                scenario: scenario_pb2.Scenario): # type: ignore
+
+        """Load from tfrecords data"""
+        states, dimensions, objects = [], [], []
+        for track in scenario.tracks:  # n_object
+            # Iterate over a single object's states.
+            track_states, track_dimensions = [], []
+            for state in track.states:  # n_timestep
+                track_states.append((state.center_x, state.center_y, state.center_z,
+                                     state.heading, state.valid))
+                track_dimensions.append((state.length, state.width, state.height))
+            # Adds to the global states.
+            states.append(list(zip(*track_states)))
+            dimensions.append(list(zip(*track_dimensions)))
+            objects.append((track.id, track.object_type))
+
+        # Unpack and convert to tf tensors.
+        x, y, z, heading, valid = [torch.tensor(s) for s in zip(*states)]
+        length, width, height = [torch.tensor(s) for s in zip(*dimensions)]
+        object_ids, object_types = [torch.tensor(s) for s in zip(*objects)]
+
+        av_id = object_ids[scenario.sdc_track_index]
+
+        init_dict = dict(x=x, y=y, z=z,
+                         heading=heading,
+                         length=length,
+                         width=width,
+                         height=height,
+                         valid=valid,
+                         object_id=object_ids,
+                         object_type=object_types,
+                         av_id=int(av_id))
+
+        return init_dict
+
+    @classmethod
+    def from_scenario(cls,
+                        scenario: scenario_pb2.Scenario, # type: ignore
+                        processed_scenario: Optional[dict]=None,
+                        from_where: str='raw'):
+
+        if from_where == 'raw':
+            init_dict = cls._get_init_dict_from_raw(scenario)
+        elif from_where == 'processed':
+            assert processed_scenario is not None, f'`processed_scenario` should be given!'
+            init_dict = cls._get_init_dict_from_processed(processed_scenario)
+        else:
+            raise RuntimeError(f'Invalid from {from_where}')
+
+        if processed_scenario is not None:
+            init_dict.update(state=processed_scenario['agent']['state_idx'],
+                             token_pos=processed_scenario['agent']['token_pos'],
+                             token_heading=processed_scenario['agent']['token_heading'],
+                             token_valid=processed_scenario['agent']['raw_agent_valid_mask'],
+                             processed_object_id=processed_scenario['agent']['id'],
+                             processed_av_id=int(processed_scenario['agent']['id'][
+                                 processed_scenario['agent']['av_idx']
+                             ]),
+            )
+
+        return cls(**init_dict)
+
+
+@dataclass
+class ScenarioRollouts:
+    scenario_id: Optional[str] = None
+    joint_scenes: List[ObjectTrajectories] = field(default_factory=list)
+
+
+""" Conversion Methods """
+
+def scenario_to_trajectories(
+    scenario: scenario_pb2.Scenario, # type: ignore
+    processed_scenario: Optional[dict]=None,
+    from_where: Optional[str]='raw',
+    remove_history: Optional[bool]=False
+) -> ObjectTrajectories:
+    """Converts a WOMD Scenario proto into the `ObjectTrajectories`.
+
+    Returns:
+    A `ObjectTrajectories` with trajectories copied from data.
+    """
+    trajectories = ObjectTrajectories.from_scenario(scenario,
+                                                    processed_scenario,
+                                                    from_where,
+                    )
+    # Slice by the required sim agents.
+    sim_agent_ids = get_sim_agent_ids(scenario)
+    # CONSOLE.log(f'sim_agent_ids of log scenario: {sim_agent_ids} total: {len(sim_agent_ids)}')
+    trajectories = trajectories.gather_objects_by_id(torch.tensor(sim_agent_ids))
+
+    if remove_history:
+        # Slice in time to only include steps after `current_time_index`.
+        trajectories = trajectories.slice_time(submission_specs.CURRENT_TIME_INDEX + 1)  # 10 + 1
+        if trajectories.valid.shape[-1] != submission_specs.N_SIMULATION_STEPS:  # 80 simulated steps
+            raise ValueError(
+                'The Scenario used does not include the right number of time steps. '
+                f'Expected: {submission_specs.N_SIMULATION_STEPS}, '
+                f'Actual: {trajectories.valid.shape[-1]}.')
+
+    return trajectories
+
+
+def _unbatch(src: Tensor, batch: Tensor, dim: int = 0) -> List[Tensor]:
+    sizes = degree(batch, dtype=torch.long).tolist()
+    return src.split(sizes, dim)
+
+
+def get_scenario_id_int_tensor(scenario_id: List[str], device: torch.device=torch.device('cpu')) -> torch.Tensor:
+    scenario_id_int_tensor = []
+    for str_id in scenario_id:
+        int_id = [-1] * 16  # max_len of scenario_id string is 16
+        for i, c in enumerate(str_id):
+            int_id[i] = ord(c)
+        scenario_id_int_tensor.append(
+            torch.tensor(int_id, dtype=torch.int32, device=device)
+        )
+    return torch.stack(scenario_id_int_tensor, dim=0)  # [n_scenario, 16]
+
+
+def output_to_rollouts(scenario: dict) -> List[ScenarioRollouts]:  # n_scenario
+    # scenario_id: Tensor,    # [n_scenario, n_str_length]
+    # agent_id: Tensor,       # [n_agent, n_rollout]
+    # agent_batch: Tensor,    # [n_agent]
+    # pred_traj: Tensor,      # [n_agent, n_rollout, n_step, 2]
+    # pred_z: Tensor,         # [n_agent, n_rollout, n_step]
+    # pred_head: Tensor,      # [n_agent, n_rollout, n_step]
+    # pred_shape: Tensor,     # [n_agent, n_rollout, 3]
+    # pred_type: Tensor,      # [n_agent, n_rollout]
+    # pred_state: Tensor,     # [n_agent, n_rollout, n_step]
+    scenario_id = scenario['scenario_id']
+    av_id = (
+        scenario['av_id'] if 'av_id' in scenario else -1
+    )
+    agent_id = scenario['agent_id']
+    agent_batch = scenario['agent_batch']
+    pred_traj = scenario['pred_traj']
+    pred_z = scenario['pred_z']
+    pred_head = scenario['pred_head']
+    pred_shape = scenario['pred_shape']
+    pred_type = scenario['pred_type']
+    pred_state = (
+        scenario['pred_state'] if 'pred_state' in scenario else
+        torch.zeros_like(pred_z).long()
+    )
+    pred_valid = scenario['pred_valid']
+    token_pos = scenario['token_pos']
+    token_head = scenario['token_head']
+
+    # CONSOLE.log("Generate scenario rollouts ...")
+    # CONSOLE.log(f'scenario_id: {scenario_id}')
+    # CONSOLE.log(f'agent_id: {agent_id.flatten()} total: {agent_id.shape}')
+    # CONSOLE.log(f'av_id: {av_id}')
+    # CONSOLE.log(f'agent_batch: {agent_batch} total: {agent_batch.shape}')
+    # CONSOLE.log(f'pred_traj: {pred_traj.shape}')
+    # CONSOLE.log(f'pred_z: {pred_z.shape}')
+    # CONSOLE.log(f'pred_head: {pred_head.shape}')
+    # CONSOLE.log(f'pred_shape: {pred_shape.shape}')
+    # CONSOLE.log(f'pred_type: {pred_type.shape}')
+    # CONSOLE.log(f'pred_state: {pred_state.shape}')
+    # CONSOLE.log(f'token_pos: {token_pos.shape}')
+    # CONSOLE.log(f'token_head: {token_head.shape}')
+
+    scenario_id = scenario_id.cpu().numpy()
+    n_agent, n_rollout, n_step, _ = pred_traj.shape
+    agent_id = _unbatch(agent_id, agent_batch)
+    pred_traj = _unbatch(pred_traj, agent_batch)
+    pred_z = _unbatch(pred_z, agent_batch)
+    pred_head = _unbatch(pred_head, agent_batch)
+    pred_shape = _unbatch(pred_shape, agent_batch)
+    pred_type = _unbatch(pred_type, agent_batch)
+    pred_state = _unbatch(pred_state, agent_batch)
+    pred_valid = _unbatch(pred_valid, agent_batch)
+    token_pos = _unbatch(token_pos, agent_batch)
+    token_head = _unbatch(token_head, agent_batch)
+
+    agent_id = [x.cpu() for x in agent_id]
+    pred_traj = [x.cpu() for x in pred_traj]
+    pred_z = [x.cpu() for x in pred_z]
+    pred_head = [x.cpu() for x in pred_head]
+    pred_shape = [x[:, :, None].repeat(1, 1, n_step, 1).cpu() for x in pred_shape]
+    pred_type = [x[:, :, None].repeat(1, 1, n_step, 1).cpu() for x in pred_type]
+    pred_state = [x.cpu() for x in pred_state]
+    pred_valid = [x.cpu() for x in pred_valid]
+    token_pos = [x.cpu() for x in token_pos]
+    token_head = [x.cpu() for x in token_head]
+
+    n_scenario = scenario_id.shape[0]
+    scenario_rollouts = []
+    for i_scenario in range(n_scenario):
+        joint_scenes = []
+        for i_rollout in range(n_rollout):  # 1
+            joint_scenes.append(
+                ObjectTrajectories(
+                    x=pred_traj[i_scenario][:, i_rollout, :, 0],
+                    y=pred_traj[i_scenario][:, i_rollout, :, 1],
+                    z=pred_z[i_scenario][:, i_rollout],
+                    heading=pred_head[i_scenario][:, i_rollout],
+                    length=pred_shape[i_scenario][:, i_rollout, :, 0],
+                    width=pred_shape[i_scenario][:, i_rollout, :, 1],
+                    height=pred_shape[i_scenario][:, i_rollout, :, 2],
+                    valid=pred_valid[i_scenario][:, i_rollout],
+                    state=pred_state[i_scenario][:, i_rollout],
+                    object_id=agent_id[i_scenario][:, i_rollout],
+                    processed_object_id=agent_id[i_scenario][:, i_rollout],
+                    object_type=pred_type[i_scenario][:, i_rollout],
+                    token_pos=token_pos[i_scenario][:, i_rollout, :, :2],
+                    token_heading=token_head[i_scenario][:, i_rollout],
+                    av_id=av_id,
+                    processed_av_id=av_id,
+                )
+            )
+
+        _str_scenario_id = "".join([chr(x) for x in scenario_id[i_scenario] if x > 0])
+        scenario_rollouts.append(
+            ScenarioRollouts(
+                joint_scenes=joint_scenes, scenario_id=_str_scenario_id
+            )
+        )
+
+    # CONSOLE.log(f'n_scenario: {len(scenario_rollouts)}')
+    # CONSOLE.log(f'n_rollout: {len(scenario_rollouts[0].joint_scenes)}')
+    # CONSOLE.log(f'x shape: {scenario_rollouts[0].joint_scenes[0].x.shape}')
+
+    return scenario_rollouts
+
+
+""" Compute Metric Features """
+
+def _compute_metametric(
+    config: long_metrics_pb2.SimAgentMetricsConfig, # type: ignore
+    metrics: long_metrics_pb2.SimAgentMetrics, # type: ignore
+):
+    """Computes the meta-metric aggregation."""
+    metametric = 0.0
+    for field_name in _METRIC_FIELD_NAMES:
+        likelihood_field_name = field_name + '_likelihood'
+        weight = getattr(config, field_name).metametric_weight
+        metric_score = getattr(metrics, likelihood_field_name)
+        metametric += weight * metric_score
+    return metametric
+
+
+@dataclasses.dataclass(frozen=True)
+class MetricFeatures:
+
+    object_id: Tensor
+    valid: Tensor
+    linear_speed: Tensor
+    linear_acceleration: Tensor
+    angular_speed: Tensor
+    angular_acceleration: Tensor
+    distance_to_nearest_object: Tensor
+    collision_per_step: Tensor
+    time_to_collision: Tensor
+    distance_to_road_edge: Tensor
+    offroad_per_step: Tensor
+    num_placement: Tensor
+    num_removement: Tensor
+    distance_placement: Tensor
+    distance_removement: Tensor
+
+    @classmethod
+    def from_file(cls, file_path: str):
+
+        if not os.path.exists(file_path):
+            raise FileNotFoundError(f'Not found file {file_path}')
+
+        with open(file_path, 'rb') as f:
+            feat_dict = pickle.load(f)
+
+        fields = [field.name for field in dataclasses.fields(cls)]
+        init_dict = dict()
+        
+        for field in fields:
+            if field in feat_dict:
+                init_dict[field] = feat_dict[field]
+            else:
+                init_dict[field] = None
+
+        return cls(**init_dict)
+
+    def unfold(self, size: int, step: int):
+        return MetricFeatures(
+            object_id=self.object_id,
+            valid=self.valid.unfold(1, size, step),
+            linear_speed=self.linear_speed.unfold(1, size, step),
+            linear_acceleration=self.linear_acceleration.unfold(1, size, step),
+            angular_speed=self.angular_speed.unfold(1, size, step),
+            angular_acceleration=self.angular_acceleration.unfold(1, size, step),
+            distance_to_nearest_object=self.distance_to_nearest_object.unfold(1, size, step),
+            collision_per_step=self.collision_per_step.unfold(1, size, step),
+            time_to_collision=self.time_to_collision.unfold(1, size, step),
+            distance_to_road_edge=self.distance_to_road_edge.unfold(1, size, step),
+            offroad_per_step=self.offroad_per_step.unfold(1, size, step),
+            num_placement=self.num_placement.unfold(1, size // SHIFT, step // SHIFT),
+            num_removement=self.num_removement.unfold(1, size // SHIFT, step // SHIFT),
+            distance_placement=self.distance_placement.unfold(1, size // SHIFT, step // SHIFT),
+            distance_removement=self.distance_removement.unfold(1, size // SHIFT, step // SHIFT),
+        )
+
+
+def compute_metric_features(
+    simulate_trajectories: ObjectTrajectories,
+    evaluate_agent_ids: Optional[Tensor]=None,
+    scenario_log: Optional[scenario_pb2.Scenario]=None, # type: ignore
+) -> MetricFeatures:
+
+    if evaluate_agent_ids is not None:
+        evaluate_trajectories = simulate_trajectories.gather_objects_by_id(
+            evaluate_agent_ids
+        )
+    else:
+        evaluate_trajectories = simulate_trajectories
+
+    # valid mask
+    validity_mask = evaluate_trajectories.valid
+    validity_mask = validity_mask[:, submission_specs.CURRENT_TIME_INDEX + 1:]
+
+    # ! Kinematics-related features, i.e. speed and acceleration, this needs
+    # history steps to be prepended to make the first evaluate step valid.
+    # Resulted `lienar_speed` and others: (n_object_to_evaluate, n_future_step)
+    linear_speed, linear_accel, angular_speed, angular_accel = (
+        trajectory_features.compute_kinematic_features(
+            evaluate_trajectories.x,
+            evaluate_trajectories.y,
+            evaluate_trajectories.z,
+            evaluate_trajectories.heading,
+            seconds_per_step=submission_specs.STEP_DURATION_SECONDS))
+    # Removes the data corresponding to the history time interval.
+    linear_speed, linear_accel, angular_speed, angular_accel = (
+        map(lambda t: t[:, submission_specs.CURRENT_TIME_INDEX + 1:],
+            [linear_speed, linear_accel, angular_speed, angular_accel])
+    )
+
+    # ! Distances to nearest objects.
+    # evaluate_object_mask = torch.any(
+    #     evaluate_agent_ids[:, None] == simulated_trajectories.object_id, axis=0
+    # )
+    evaluate_object_mask = torch.ones(len(simulate_trajectories.object_id)).bool()
+    distances_to_objects = interact_features.compute_distance_to_nearest_object(
+        center_x=simulate_trajectories.x,
+        center_y=simulate_trajectories.y,
+        center_z=simulate_trajectories.z,
+        length=simulate_trajectories.length,
+        width=simulate_trajectories.width,
+        height=simulate_trajectories.height,
+        heading=simulate_trajectories.heading,
+        valid=simulate_trajectories.valid,
+        evaluated_object_mask=evaluate_object_mask,
+    )
+    distances_to_objects = (
+        distances_to_objects[:, submission_specs.CURRENT_TIME_INDEX + 1:])
+    is_colliding_per_step = torch.lt(
+        distances_to_objects, interact_features.COLLISION_DISTANCE_THRESHOLD)
+
+    # ! Time to collision
+    times_to_collision = (
+        interact_features.compute_time_to_collision_with_object_in_front(
+            center_x=simulate_trajectories.x,
+            center_y=simulate_trajectories.y,
+            length=simulate_trajectories.length,
+            width=simulate_trajectories.width,
+            heading=simulate_trajectories.heading,
+            valid=simulate_trajectories.valid,
+            evaluated_object_mask=evaluate_object_mask,
+            seconds_per_step=submission_specs.STEP_DURATION_SECONDS,
+        )
+    )
+    times_to_collision = times_to_collision[:, submission_specs.CURRENT_TIME_INDEX + 1:]
+
+    # ! Distance to road edge
+    distances_to_road_edge = torch.empty_like(distances_to_objects)
+    is_offroad_per_step = torch.empty_like(is_colliding_per_step)
+    if scenario_log is not None:
+        road_edges = []
+        for map_feature in scenario_log.map_features:
+            if map_feature.HasField('road_edge'):
+                road_edges.append(map_feature.road_edge.polyline)
+        distances_to_road_edge = map_features.compute_distance_to_road_edge(
+            center_x=simulate_trajectories.x,
+            center_y=simulate_trajectories.y,
+            center_z=simulate_trajectories.z,
+            length=simulate_trajectories.length,
+            width=simulate_trajectories.width,
+            height=simulate_trajectories.height,
+            heading=simulate_trajectories.heading,
+            valid=simulate_trajectories.valid,
+            evaluated_object_mask=evaluate_object_mask,
+            road_edge_polylines=road_edges,
+        )
+        distances_to_road_edge = distances_to_road_edge[:, submission_specs.CURRENT_TIME_INDEX + 1:]
+        is_offroad_per_step = torch.gt(
+            distances_to_road_edge, map_features.OFFROAD_DISTANCE_THRESHOLD
+        )
+
+    # ! Placement
+    if simulate_trajectories.av_id == simulate_trajectories.processed_av_id == -1:
+        n_agent, n_step_10hz = linear_speed.shape
+        num_placement = torch.zeros((n_step_10hz // SHIFT,))
+        num_removement = torch.zeros((n_step_10hz // SHIFT,))
+        distance_placement = torch.zeros((n_agent, n_step_10hz // SHIFT))
+        distance_removement = torch.zeros((n_agent, n_step_10hz // SHIFT))
+
+    else:
+        assert simulate_trajectories.av_id == simulate_trajectories.processed_av_id, \
+                f"Got duplicated av_id: {simulate_trajectories.av_id} and {simulate_trajectories.processed_av_id}"
+        num_placement, num_removement = (
+            placement_features.compute_num_placement(
+                state=simulate_trajectories.state,
+                valid=simulate_trajectories.token_valid,
+                av_id=simulate_trajectories.processed_av_id,
+                object_id=simulate_trajectories.processed_object_id,
+                agent_state=AGENT_STATE,
+            )
+        )
+        num_placement = num_placement[submission_specs.CURRENT_TIME_INDEX // SHIFT:]
+        num_removement = num_removement[submission_specs.CURRENT_TIME_INDEX // SHIFT:]
+        distance_placement, distance_removement = (
+            placement_features.compute_distance_placement(
+                position=simulate_trajectories.token_pos,
+                state=simulate_trajectories.state,
+                valid=simulate_trajectories.valid,
+                av_id=simulate_trajectories.processed_av_id,
+                object_id=simulate_trajectories.processed_object_id,
+                agent_state=AGENT_STATE,
+            )
+        )
+        distance_placement = distance_placement[:, submission_specs.CURRENT_TIME_INDEX // SHIFT:]
+        distance_removement = distance_removement[:, submission_specs.CURRENT_TIME_INDEX // SHIFT:]
+        # distance_placement = distance_placement[distance_placement > 0]
+        # distance_removement = distance_removement[distance_removement > 0]
+
+        # print out some results for debugging
+        # CONSOLE.log(f'trajectory x: {simulate_trajectories.x.shape}, \n{simulate_trajectories.x}')
+        # CONSOLE.log(f'linear speed: {linear_speed.shape}, \n{linear_speed}')
+        # CONSOLE.log(f'distances: {distances_to_objects.shape}, \n{distances_to_objects}')
+        # CONSOLE.log(f'time to collision: {times_to_collision.shape}, {times_to_collision}')
+
+    return MetricFeatures(
+            object_id=simulate_trajectories.object_id,
+            valid=validity_mask,
+            # kinematic
+            linear_speed=linear_speed,
+            linear_acceleration=linear_accel,
+            angular_speed=angular_speed,
+            angular_acceleration=angular_accel,
+            # interact
+            distance_to_nearest_object=distances_to_objects,
+            collision_per_step=is_colliding_per_step,
+            time_to_collision=times_to_collision,
+            # map
+            distance_to_road_edge=distances_to_road_edge,
+            offroad_per_step=is_offroad_per_step,
+            # placement
+            num_placement=num_placement[None, ...],
+            num_removement=num_removement[None, ...],
+            distance_placement=distance_placement,
+            distance_removement=distance_removement,
+    )
+
+
+@dataclass(frozen=True)
+class LogDistributions:
+
+    linear_speed: Tensor
+    linear_acceleration: Tensor
+    angular_speed: Tensor
+    angular_acceleration: Tensor
+    distance_to_nearest_object: Tensor
+    collision_indication: Tensor
+    time_to_collision: Tensor
+    distance_to_road_edge: Tensor
+    num_placement: Tensor
+    num_removement: Tensor
+    distance_placement: Tensor
+    distance_removement: Tensor
+    offroad_indication: Optional[Tensor] = None
+
+
+""" Compute Metrics """
+
+def _assert_and_return_batch_size(
+    log_samples: Tensor, 
+    sim_samples: Tensor
+) -> int:
+    """Asserts consistency in the tensor shapes and returns batch size.
+
+    Args:
+        log_samples: A tensor of shape (batch_size, log_sample_size).
+        sim_samples: A tensor of shape (batch_size, sim_sample_size).
+
+    Returns:
+        The `batch_size`.
+    """
+    assert log_samples.shape[0] == sim_samples.shape[0], "Log and Sim batch sizes must be equal."
+    return log_samples.shape[0]
+
+
+def _reduce_average_with_validity(
+    tensor: Tensor, validity: Tensor) -> Tensor:
+    """Returns the tensor's average, only selecting valid items.
+
+    Args:
+    tensor: A float tensor of any shape.
+    validity: A boolean tensor of the same shape as `tensor`.
+
+    Returns:
+    A float tensor of shape (1,), containing the average of the valid elements
+    of `tensor`.
+    """
+    if tensor.shape != validity.shape:
+        raise ValueError('Shapes of `tensor` and `validity` must be the same.'
+                            f'(Actual: {tensor.shape}, {validity.shape}).')
+    cond_sum = torch.sum(torch.where(validity, tensor, torch.zeros_like(tensor)))
+    valid_sum = torch.sum(validity)
+    if valid_sum == 0:
+        return torch.tensor(0.)
+    return cond_sum / valid_sum
+
+
+def histogram_estimate(
+    config: long_metrics_pb2.SimAgentMetricsConfig.HistogramEstimate, # type: ignore
+    log_samples: Tensor,
+    sim_samples: Tensor,
+) -> Tensor:
+    """Computes log-likelihoods of samples based on histograms.
+
+    Args:
+        config: A configuration dictionary, similar to the one in TensorFlow.
+        log_samples: A tensor of shape (batch_size, log_sample_size),
+            containing `log_sample_size` samples from `batch_size` independent
+            populations.
+        sim_samples: A tensor of shape (batch_size, sim_sample_size),
+            containing `sim_sample_size` samples from `batch_size` independent
+            populations.
+
+    Returns:
+        A tensor of shape (batch_size, log_sample_size), where each element (i, k)
+        is the log likelihood of the log sample (i, k) under the sim distribution
+        (i).
+    """
+    batch_size = _assert_and_return_batch_size(log_samples, sim_samples)
+
+    # We generate `num_bins`+1 edges for the histogram buckets.
+    edges = torch.linspace(
+        config.min_val, config.max_val, config.num_bins + 1
+    ).float()
+
+    # Clip the samples to avoid errors with histograms.
+    log_samples = torch.clamp(log_samples, config.min_val, config.max_val)
+    sim_samples = torch.clamp(sim_samples, config.min_val, config.max_val)
+
+    # Create the categorical distribution for simulation. `tfp.histogram` returns
+    # a tensor of shape (num_bins, batch_size), so we need to transpose to conform
+    # to `tfp.distribution.Categorical`, which requires `probs` to be
+    # (batch_size, num_bins).
+    sim_counts = torch.vmap(lambda x: torch.histogram(x, bins=edges).hist)(sim_samples)
+    sim_counts += config.additive_smoothing_pseudocount
+    distributions = torch.distributions.Categorical(probs=sim_counts)
+
+    # Generate the counts for the log distribution. We reshape the log samples to
+    # (batch_size * log_sample_size, 1), so every log sample is independently
+    # scored.
+    log_values_flat = log_samples.reshape(-1, 1)
+    # Shape of log_counts: (batch_size * log_sample_size, num_bins).
+    log_counts = torch.vmap(lambda x: torch.histogram(x, bins=edges).hist)(log_values_flat)
+    # Identify which bin each sample belongs to and get the log probability of
+    # that bin under the sim distribution.
+    max_log_bin = log_counts.argmax(dim=-1)
+    batched_max_log_bin = max_log_bin.reshape(batch_size, -1)
+
+    # Since we have defined the categorical distribution to have `batch_size`
+    # independent populations, tfp expects this `batch_size` to be in the last
+    # dimension of the tensor, so transpose the log bins to
+    # (log_sample_size, batch_size).
+    log_likelihood = distributions.log_prob(batched_max_log_bin.transpose(0, 1))
+
+    # Return log likelihood in the shape (batch_size, log_sample_size)
+    return log_likelihood.transpose(0, 1)
+
+
+def log_likelihood_estimate_timeseries(
+    field: str,
+    feature_config: long_metrics_pb2.SimAgentMetricsConfig.FeatureConfig, # type: ignore
+    sim_values: Tensor,
+    log_distributions: torch.distributions.Categorical,
+    estimate_method: str='histogram',
+) -> Tensor:
+    """Computes the log-likelihood estimates for a time-series simulated feature.
+
+    Args:
+    feature_config: A time-series compatible `FeatureConfig`.
+    log_distributions: A float Tensor with shape (batch_sizie, n_bins).
+    sim_values: A float Tensor with shape (n_objects / n_scenarios, n_segments, n_steps).
+
+    Returns:
+    A tensor of shape (n_objects, n_steps) containing the simulation probability
+    estimates of the simulation features under the logged distribution of the same
+    feature.
+    """
+    assert sim_values.ndim == 3, f'Expect sim_values.ndim==3, got {sim_values.ndim}, shape {sim_values.shape} for {field}' 
+
+    sim_values_flat = sim_values.reshape(-1, 1)  # [n_objects * n_segments * n_steps]
+
+    # if not feature_config.independent_timesteps:
+    #     # If time steps needs to be considered independent, reshape:
+    #     # - `sim_values` as (n_objects, n_rollouts * n_steps)
+    #     # - `log_values` as (n_objects, n_steps)
+    #     # If values in time are instead to be compared per-step, reshape:
+    #     # - `sim_values` as (n_objects * n_steps, n_rollouts)
+    #     # - `log_values` as (n_objects * n_steps, 1)
+    #     sim_values = sim_values.reshape(-1, 1)  # n_rollouts=1
+
+    # if feature_config.independent_timesteps:
+    #     sim_values = sim_values.permute(1, 0, 2).reshape(n_objects, n_rollouts * n_steps)
+    # else:
+    #     sim_values = sim_values.permute(1, 2, 0).reshape(n_objects * n_steps, n_rollouts)
+    #     log_values = log_values.reshape(n_objects * n_steps, 1)
+
+    # ! calculate distributions for simulate features
+    if estimate_method == 'histogram':
+        config = feature_config.histogram
+    elif estimate_method == 'bernoulli':
+        config = (
+            long_metrics_pb2.SimAgentMetricsConfig.HistogramEstimate(
+                min_val=-0.5, max_val=0.5, num_bins=2,
+                additive_smoothing_pseudocount=feature_config.bernoulli.additive_smoothing_pseudocount
+            )
+        )
+        sim_values_flat = sim_values_flat.float()  # cast torch.bool to torch.float32
+
+    # We generate `num_bins`+1 edges for the histogram buckets.
+    edges = torch.linspace(
+        config.min_val, config.max_val, config.num_bins + 1
+    ).float()
+
+    sim_counts = torch.vmap(lambda x: torch.histogram(x, bins=edges).hist)(sim_values_flat)  # [batch_size, num_bins]
+    # Identify which bin each sample belongs to and get the log probability of
+    # that bin under the sim distribution.
+    max_sim_bin = sim_counts.argmax(dim=-1)
+    batched_max_sim_bin = max_sim_bin.reshape(1, -1)  # `batch_size` = 1, follows the log distributions
+
+    sim_likelihood = log_distributions.log_prob(batched_max_sim_bin.transpose(0, 1)).flatten()
+    return sim_likelihood.reshape(*sim_values.shape)  # [n_objects, n_segments, n_steps]
+
+
+def compute_scenario_metrics_for_bundle(
+    config: long_metrics_pb2.SimAgentMetricsConfig, # type: ignore
+    log_distributions: LogDistributions,
+    scenario_log: Optional[scenario_pb2.Scenario], # type: ignore
+    scenario_rollouts: ScenarioRollouts,
+) -> long_metrics_pb2.SimAgentMetrics: # type: ignore
+
+    features_fields = [field.name for field in dataclasses.fields(MetricFeatures)]
+    features_fields.remove('object_id')
+
+    # ! compute simluation features
+    # CONSOLE.log('[on yellow] Compute sim features [/]')
+    sim_features = collections.defaultdict(list)
+    for simulate_trajectories in tqdm(scenario_rollouts.joint_scenes, leave=False, desc='rollouts ...'):  # n_rollout=1
+        rollout_features = compute_metric_features(
+            simulate_trajectories,
+            evaluate_agent_ids=None,
+            scenario_log=scenario_log
+        )
+
+        for field in features_fields:
+            sim_features[field].append(getattr(rollout_features, field))
+
+    for field in features_fields:
+        if sim_features[field][0] is not None:
+            sim_features[field] = torch.concat(sim_features[field], dim=0)  # n_rollout for dim=0
+
+    sim_features = MetricFeatures(
+        **sim_features, object_id=None,
+    )
+    # after unfold: linear_speed shape [n_agent, n_window, window_size],
+    # num_placement shape [n_scenario=1, n_window, window_size]
+    flattened_sim_features = copy.deepcopy(sim_features)
+    sim_features = sim_features.unfold(size=submission_specs.N_SIMULATION_STEPS, step=SHIFT)
+    # CONSOLE.log(f'sim linear_speed feature: {sim_features.linear_speed.shape}')
+    # CONSOLE.log(f'sim num_placement feature: {sim_features.num_placement.shape}')
+
+    ## ! compute metrics
+
+    # ! kinematics-related metrics
+    linear_speed_log_likelihood = log_likelihood_estimate_timeseries(
+        field='linear_speed',
+        feature_config=config.linear_speed,
+        sim_values=sim_features.linear_speed,
+        log_distributions=log_distributions.linear_speed,
+    )
+    angular_speed_log_likelihood = log_likelihood_estimate_timeseries(
+        field='angular_speed',
+        feature_config=config.angular_speed,
+        sim_values=sim_features.angular_speed,
+        log_distributions=log_distributions.angular_speed,
+    )
+    speed_validity, acceleration_validity = (
+        trajectory_features.compute_kinematic_validity(flattened_sim_features.valid)
+    )
+    speed_validity = speed_validity.unfold(1, size=submission_specs.N_SIMULATION_STEPS, step=SHIFT)
+    acceleration_validity = acceleration_validity.unfold(1, size=submission_specs.N_SIMULATION_STEPS, step=SHIFT)
+    linear_speed_likelihood = torch.exp(_reduce_average_with_validity(
+        linear_speed_log_likelihood, speed_validity))
+    angular_speed_likelihood = torch.exp(_reduce_average_with_validity(
+        angular_speed_log_likelihood, speed_validity))
+    # CONSOLE.log(f'linear_speed_likelihood: {linear_speed_likelihood}')
+    # CONSOLE.log(f'angular_speed_likelihood: {angular_speed_likelihood}')
+
+    linear_accel_log_likelihood = log_likelihood_estimate_timeseries(
+        field='linear_acceleration',
+        feature_config=config.linear_acceleration,
+        sim_values=sim_features.linear_acceleration,
+        log_distributions=log_distributions.linear_acceleration,
+    )
+    angular_accel_log_likelihood = log_likelihood_estimate_timeseries(
+        field='angular_acceleration',
+        feature_config=config.angular_acceleration,
+        sim_values=sim_features.angular_acceleration,
+        log_distributions=log_distributions.angular_acceleration,
+    )
+    linear_accel_likelihood = torch.exp(_reduce_average_with_validity(
+        linear_accel_log_likelihood, acceleration_validity))
+    angular_accel_likelihood = torch.exp(_reduce_average_with_validity(
+        angular_accel_log_likelihood, acceleration_validity))
+    # CONSOLE.log(f'linear_accel_likelihood: {linear_accel_likelihood}')
+    # CONSOLE.log(f'angular_accel_likelihood: {angular_accel_likelihood}')
+
+    # ! collision and distance to other objects.
+
+    sim_collision_indication = torch.any(
+        torch.where(sim_features.valid, sim_features.collision_per_step, False),
+        dim=2)[..., None]  # add a dummy time dimension
+    collision_score = log_likelihood_estimate_timeseries(
+        field='collision_indication',
+        feature_config=config.collision_indication,
+        sim_values=sim_collision_indication,
+        log_distributions=log_distributions.collision_indication,
+        estimate_method='bernoulli',
+    )
+    collision_likelihood = torch.exp(torch.mean(collision_score))
+
+    distance_to_objects_log_likelihodd = log_likelihood_estimate_timeseries(
+        field='distance_to_nearest_object',
+        feature_config=config.distance_to_nearest_object,
+        sim_values=sim_features.distance_to_nearest_object,
+        log_distributions=log_distributions.distance_to_nearest_object,
+    )
+    distance_to_objects_likelihodd = torch.exp(_reduce_average_with_validity(
+        distance_to_objects_log_likelihodd, sim_features.valid))
+    # CONSOLE.log(f'distance_to_objects_likelihodd: {distance_to_objects_likelihodd}')
+
+    ttc_log_likelihood = log_likelihood_estimate_timeseries(
+        field='time_to_collision',
+        feature_config=config.time_to_collision,
+        sim_values=sim_features.time_to_collision,
+        log_distributions=log_distributions.time_to_collision,
+    )
+    ttc_likelihood = torch.exp(_reduce_average_with_validity(
+        ttc_log_likelihood, sim_features.valid))
+    # CONSOLE.log(f'ttc_likelihood: {ttc_likelihood}')
+
+    # ! offroad and distance to road edge.
+
+    # distance_to_road_edge_log_likelihood = log_likelihood_estimate_timeseries(
+    #     field='distance_to_road_edge',
+    #     sim_values=sim_features.distance_to_road_edge,
+    #     log_distributions=log_distributions.distance_to_road_edge,
+    # )
+    # distance_to_road_edge_likelihood = torch.exp(_reduce_average_with_validity(
+    #     distance_to_road_edge_log_likelihood, sim_features.valid))
+    # CONSOLE.log(f'distance_to_road_edge_likelihood: {distance_to_road_edge_likelihood}')
+
+    # ! placement
+
+    num_placement_log_likelihood = log_likelihood_estimate_timeseries(
+        field='num_placement',
+        feature_config=config.num_placement,
+        sim_values=sim_features.num_placement.float(),
+        log_distributions=log_distributions.num_placement,
+    )
+    num_placement_likelihood = torch.exp(torch.mean(num_placement_log_likelihood))
+    num_removement_log_likelihood = log_likelihood_estimate_timeseries(
+        field='num_removement',
+        feature_config=config.num_removement,
+        sim_values=sim_features.num_removement.float(),
+        log_distributions=log_distributions.num_removement,
+    )
+    num_removement_likelihood = torch.exp(torch.mean(num_removement_log_likelihood))
+    # CONSOLE.log(f'num_placement_likelihood: {num_placement_likelihood}')
+    # CONSOLE.log(f'num_removement_likelihood: {num_removement_likelihood}')
+
+    # tensor([[0.0013, 0.0078, 0.0194, 0.0373, 0.0628, 0.0938, 0.1232, 0.1470, 0.1701,
+    #         0.3371]])
+    # tensor([[0.0201, 0.0570, 0.0689, 0.0839, 0.1029, 0.1172, 0.1282, 0.1286, 0.1237,
+    #         0.1695]])
+    distance_placement_log_likelihood = log_likelihood_estimate_timeseries(
+        field='distance_placement',
+        feature_config=config.distance_placement,
+        sim_values=sim_features.distance_placement,
+        log_distributions=log_distributions.distance_placement,
+    )
+    distance_placement_validity = (
+        (sim_features.distance_placement > config.distance_placement.histogram.min_val) &
+        (sim_features.distance_placement < config.distance_placement.histogram.max_val)
+    )
+    distance_placement_likelihood = torch.exp(_reduce_average_with_validity(
+        distance_placement_log_likelihood, distance_placement_validity))
+    distance_removement_log_likelihood = log_likelihood_estimate_timeseries(
+        field='distance_removement',
+        feature_config=config.distance_removement,
+        sim_values=sim_features.distance_removement,
+        log_distributions=log_distributions.distance_removement,
+    )
+    distance_removement_validity = (
+        (sim_features.distance_removement > config.distance_removement.histogram.min_val) &
+        (sim_features.distance_removement < config.distance_removement.histogram.max_val)
+    )
+    distance_removement_likelihood = torch.exp(_reduce_average_with_validity(
+        distance_removement_log_likelihood, distance_removement_validity))
+
+    # ==== Simulated collision and offroad rates ====
+    simulated_collision_rate = torch.sum(
+        sim_collision_indication.long()
+    ) / torch.sum(torch.ones_like(sim_collision_indication).long())
+    # simulated_offroad_rate = tf.reduce_sum(
+    #     # `sim_offroad_indication` shape: (n_samples, n_objects).
+    #     tf.cast(sim_offroad_indication, tf.int32)
+    # ) / tf.reduce_sum(tf.ones_like(sim_offroad_indication, dtype=tf.int32))
+
+    # ==== Meta metric ====
+    likelihood_metrics = {
+        'linear_speed_likelihood': float(linear_speed_likelihood.numpy()),
+        'linear_acceleration_likelihood': float(linear_accel_likelihood.numpy()),
+        'angular_speed_likelihood': float(angular_speed_likelihood.numpy()),
+        'angular_acceleration_likelihood': float(angular_accel_likelihood.numpy()),
+        'distance_to_nearest_object_likelihood': float(distance_to_objects_likelihodd.numpy()),
+        'collision_indication_likelihood': float(collision_likelihood.numpy()),
+        'time_to_collision_likelihood': float(ttc_likelihood.numpy()),
+        # 'distance_to_road_edge_likelihoodfloat(': distance_road_edge_likelihood.nump)y(),
+        # 'offroad_indication_likelihoodfloat(': offroad_likelihood.nump)y(),
+        'num_placement_likelihood': float(num_placement_likelihood.numpy()),
+        'num_removement_likelihood': float(num_removement_likelihood.numpy()),
+        'distance_placement_likelihood': float(distance_placement_likelihood.numpy()),
+        'distance_removement_likelihood': float(distance_removement_likelihood.numpy()),
+    }
+
+    metametric = _compute_metametric(
+        config, long_metrics_pb2.SimAgentMetrics(**likelihood_metrics)
+    )
+    # CONSOLE.log(f'metametric: {metametric}')
+
+    return long_metrics_pb2.SimAgentMetrics(
+        scenario_id=scenario_rollouts.scenario_id,
+        metametric=metametric,
+        simulated_collision_rate=float(simulated_collision_rate.numpy()),
+        # simulated_offroad_rate=simulated_offroad_rate.numpy(),
+        **likelihood_metrics,
+    )
+
+
+""" Log Features """
+
+def _get_log_distributions(
+    field: str,
+    feature_config: long_metrics_pb2.SimAgentMetricsConfig.FeatureConfig, # type: ignore
+    log_values: Tensor,
+    estimate_method: str = 'histogram',
+) -> Tensor:
+    """Computes the log-likelihood estimates for a time-series simulated feature.
+
+    Args:
+    feature_config: A time-series compatible `FeatureConfig`.
+    log_values: A float Tensor with shape (n_objects, n_steps).
+    sim_values: A float Tensor with shape (n_rollouts, n_objects, n_steps).
+
+    Returns:
+    A tensor of shape (n_objects, n_steps) containing the log probability
+    estimates of the log features under the simulated distribution of the same
+    feature.
+    """
+    assert log_values.ndim == 2, f'Expect log_values.ndim==2, got {log_values.ndim}, shape {log_values.shape} for {field}' 
+
+    # [n_objects, n_steps] -> [n_objects * n_steps]
+    log_samples = log_values.reshape(-1)
+
+    # ! estimate
+    if estimate_method == 'histogram':
+        config = feature_config.histogram
+    elif estimate_method == 'bernoulli':
+        config = (
+            long_metrics_pb2.SimAgentMetricsConfig.HistogramEstimate(
+                min_val=-0.5, max_val=0.5, num_bins=2,
+                additive_smoothing_pseudocount=feature_config.bernoulli.additive_smoothing_pseudocount
+            )
+        )
+        log_samples = log_samples.float()  # cast torch.bool to torch.float32
+
+    # We generate `num_bins`+1 edges for the histogram buckets.
+    edges = torch.linspace(
+        config.min_val, config.max_val, config.num_bins + 1
+    ).float()
+
+    if field in ('distance_placement', 'distance_removement'):
+        log_samples = log_samples[(log_samples > config.min_val) & (log_samples < config.max_val)]
+
+    # Clip the samples to avoid errors with histograms. Nonetheless, the min/max
+    # values should be configured to never hit this condition in practice.
+    log_samples = torch.clamp(log_samples, config.min_val, config.max_val)
+
+    # Create the categorical distribution for simulation. `tfp.histogram` returns
+    # a tensor of shape (num_bins, batch_size), so we need to transpose to conform
+    # to `tfp.distribution.Categorical`, which requires `probs` to be
+    # (batch_size, num_bins).
+    log_counts = torch.histogram(log_samples, bins=edges).hist.unsqueeze(dim=0)  # [1, n_samples]
+    log_counts += config.additive_smoothing_pseudocount
+    distributions = torch.distributions.Categorical(probs=log_counts)
+
+    return distributions
+
+
+class LongMetric(Metric):
+
+    log_features: MetricFeatures
+
+    def __init__(
+            self,
+            prefix: str='',
+            log_features_dir: str='data/waymo_processed/log_features/',
+            config_path: str='dev/metrics/metric_config.textproto',
+    ) -> None:
+        super().__init__()
+        self.prefix = prefix
+        self.metrics_config = self.load_metrics_config(config_path)
+
+        self.use_log = False
+
+        self.field_names = [
+            "metametric",
+            "average_displacement_error",
+            "min_average_displacement_error",
+            "linear_speed_likelihood",
+            "linear_acceleration_likelihood",
+            "angular_speed_likelihood",
+            "angular_acceleration_likelihood",
+            'distance_to_nearest_object_likelihood',
+            'collision_indication_likelihood',
+            'time_to_collision_likelihood',
+            # 'distance_to_road_edge_likelihood',
+            # 'offroad_indication_likelihood',
+            'simulated_collision_rate',
+            # 'simulated_offroad_rate',
+            'num_placement_likelihood',
+            'num_removement_likelihood',
+            'distance_placement_likelihood',
+            'distance_removement_likelihood',
+        ]
+        for k in self.field_names:
+            self.add_state(k, default=torch.tensor(0.), dist_reduce_fx='sum')
+        self.add_state('scenario_counter', default=torch.tensor(0.), dist_reduce_fx='sum')
+        self.add_state('placement_valid_scenario_counter', default=torch.tensor(0.), dist_reduce_fx='sum')
+        self.add_state('removement_valid_scenario_counter', default=torch.tensor(0.), dist_reduce_fx='sum')
+
+        # get log features
+        log_features_path = os.path.join(log_features_dir, 'total_features.pkl')
+        if not os.path.exists(log_features_path):
+            CONSOLE.log(f'[on yellow] Log features does not exist, loading now ... [/]')
+            log_features = aggregate_log_metric_features(log_features_dir)
+        else:
+            log_features = MetricFeatures.from_file(log_features_path)
+            CONSOLE.log(f'Loaded log features from {log_features_path}')
+        self.log_features = log_features
+
+        self._compute_distributions()
+        CONSOLE.log(f"Calculated log distributions:\n{self.log_distributions}")
+
+    def _compute_distributions(self):
+        self.log_distributions = LogDistributions(
+            linear_speed = _get_log_distributions('linear_speed',
+                self.metrics_config.linear_speed, self.log_features.linear_speed,
+            ),
+            linear_acceleration = _get_log_distributions('linear_acceleration',
+                self.metrics_config.linear_acceleration, self.log_features.linear_acceleration,
+            ),
+            angular_speed = _get_log_distributions('angular_speed',
+                self.metrics_config.angular_speed, self.log_features.angular_speed,
+            ),
+            angular_acceleration = _get_log_distributions('angular_acceleration',
+                self.metrics_config.angular_acceleration, self.log_features.angular_acceleration,
+            ),
+            distance_to_nearest_object = _get_log_distributions('distance_to_nearest_object',
+                self.metrics_config.distance_to_nearest_object, self.log_features.distance_to_nearest_object,
+            ),
+            collision_indication = _get_log_distributions('collision_indication',
+                self.metrics_config.collision_indication,
+                log_collision_indication := torch.any(
+                    torch.where(self.log_features.valid, self.log_features.collision_per_step, False), dim=1
+                )[..., None],  # add a dummy time dimension
+                estimate_method = 'bernoulli',
+            ),
+            time_to_collision = _get_log_distributions('time_to_collision',
+                self.metrics_config.time_to_collision, self.log_features.time_to_collision,
+            ),
+            distance_to_road_edge = _get_log_distributions('distance_to_road_edge',
+                self.metrics_config.distance_to_road_edge, self.log_features.distance_to_road_edge,
+            ),
+            # dist_offroad_indication = _get_log_distributions(
+            #     'offroad_indication',
+            #     self.metrics_config.offroad_indication,
+            #     log_offroad_indication := torch.any(
+            #         torch.where(self.log_features.valid, self.log_features.offroad_per_step, False), dim=1
+            #     ),
+            # ),
+            num_placement = _get_log_distributions('num_placement',
+                self.metrics_config.num_placement, self.log_features.num_placement.float(),
+            ),
+            num_removement = _get_log_distributions('num_removement',
+                self.metrics_config.num_removement, self.log_features.num_removement.float(),
+            ),
+            distance_placement = _get_log_distributions('distance_placement',
+                self.metrics_config.distance_placement, (
+                    self.log_features.distance_placement[self.log_features.distance_placement > 0])[None, ...],
+            ),
+            distance_removement = _get_log_distributions('distance_removement',
+                self.metrics_config.distance_removement, (
+                    self.log_features.distance_removement[self.log_features.distance_removement > 0])[None, ...],
+            ),
+        )
+
+    def _compute_scenario_metrics(
+        self,
+        scenario_file: Optional[str],
+        scenario_rollout: ScenarioRollouts,
+    ) -> long_metrics_pb2.SimAgentMetrics: # type: ignore
+
+        scenario_log = None
+        if self.use_log and scenario_file is not None:
+            if not os.path.exists(scenario_file):
+                raise FileNotFoundError(f"Not found file {scenario_file}")
+            scenario_log = scenario_pb2.Scenario()
+            for data in tf.data.TFRecordDataset([scenario_file], compression_type=''):
+                scenario_log.ParseFromString(bytes(data.numpy()))
+                break
+
+        return compute_scenario_metrics_for_bundle(
+            self.metrics_config, self.log_distributions, scenario_log, scenario_rollout
+        )
+
+    def compute_metrics(self, outputs: dict) -> List[long_metrics_pb2.SimAgentMetrics]: # type: ignore
+        """
+        `outputs` is a dict directly generated by predict models:
+            >>> outputs = dict(
+            >>>     scenario_id=get_scenario_id_int_tensor(data['scenario_id'], device),
+            >>>     agent_id=agent_id,
+            >>>     agent_batch=agent_batch,
+            >>>     pred_traj=pred_traj,
+            >>>     pred_z=pred_z,
+            >>>     pred_head=pred_head,
+            >>>     pred_shape=pred_shape,
+            >>>     pred_type=pred_type,
+            >>>     pred_state=pred_state,
+            >>> )
+        """
+
+        scenario_rollouts = output_to_rollouts(outputs)
+        log_paths: List[str] = outputs['tfrecord_path']
+
+        pool_scenario_metrics = []
+        for _scenario_file, _scenario_rollout in tqdm(
+            zip(log_paths, scenario_rollouts), leave=False, desc='scenarios ...'):  # n_scenarios
+            pool_scenario_metrics.append(
+                self._compute_scenario_metrics(
+                    _scenario_file, _scenario_rollout,
+                )
+            )
+
+        return pool_scenario_metrics
+
+    def update(
+            self,
+            outputs: Optional[dict]=None,
+            metrics: Optional[List[long_metrics_pb2.SimAgentMetrics]]=None # type: ignore
+        ) -> None:
+
+        if metrics is None:
+            assert outputs is not None, f'`outputs` should not be None!'
+            metrics = self.compute_metrics(outputs)
+
+        for scenario_metrics in metrics:
+            self.scenario_counter += 1
+
+            self.metametric += scenario_metrics.metametric
+            self.average_displacement_error += (
+                scenario_metrics.average_displacement_error
+            )
+            self.min_average_displacement_error += (
+                scenario_metrics.min_average_displacement_error
+            )
+            self.linear_speed_likelihood += scenario_metrics.linear_speed_likelihood
+            self.linear_acceleration_likelihood += (
+                scenario_metrics.linear_acceleration_likelihood
+            )
+            self.angular_speed_likelihood += scenario_metrics.angular_speed_likelihood
+            self.angular_acceleration_likelihood += (
+                scenario_metrics.angular_acceleration_likelihood
+            )
+            self.distance_to_nearest_object_likelihood += (
+                scenario_metrics.distance_to_nearest_object_likelihood
+            )
+            self.collision_indication_likelihood += (
+                scenario_metrics.collision_indication_likelihood
+            )
+            self.time_to_collision_likelihood += (
+                scenario_metrics.time_to_collision_likelihood
+            )
+            # self.distance_to_road_edge_likelihood += (
+            #     scenario_metrics.distance_to_road_edge_likelihood
+            # )
+            # self.offroad_indication_likelihood += (
+            #     scenario_metrics.offroad_indication_likelihood
+            # )
+            self.simulated_collision_rate += scenario_metrics.simulated_collision_rate
+            # self.simulated_offroad_rate += scenario_metrics.simulated_offroad_rate
+
+            self.num_placement_likelihood += (
+                scenario_metrics.num_placement_likelihood
+            )
+            self.num_removement_likelihood += (
+                scenario_metrics.num_removement_likelihood
+            )
+            self.distance_placement_likelihood += (
+                scenario_metrics.distance_placement_likelihood
+            )
+            self.distance_removement_likelihood += (
+                scenario_metrics.distance_removement_likelihood
+            )
+
+            if scenario_metrics.distance_placement_likelihood > 0:
+                self.placement_valid_scenario_counter += 1
+
+            if scenario_metrics.distance_removement_likelihood > 0:
+                self.removement_valid_scenario_counter += 1
+
+    def compute(self) -> Dict[str, Tensor]:
+        metrics_dict = {}
+        for k in self.field_names:
+            if k not in ('distance_placement', 'distance_removement'):
+                metrics_dict[k] = getattr(self, k) / self.scenario_counter
+            if k == 'distance_placement':
+                metrics_dict[k] = getattr(self, k) / self.placement_valid_scenario_counter
+            if k == 'distance_removement':
+                metrics_dict[k] = getattr(self, k) / self.removement_valid_scenario_counter
+
+        mean_metrics = long_metrics_pb2.SimAgentMetrics(
+            scenario_id='', **metrics_dict,
+        )
+        final_metrics = self.aggregate_metrics_to_buckets(
+            self.metrics_config, mean_metrics
+        )
+        CONSOLE.log(f'final_metrics:\n{final_metrics}')
+
+        out_dict = {
+            f"{self.prefix}/wosac/realism_meta_metric": final_metrics.realism_meta_metric,
+            f"{self.prefix}/wosac/kinematic_metrics": final_metrics.kinematic_metrics,
+            f"{self.prefix}/wosac/interactive_metrics": final_metrics.interactive_metrics,
+            f"{self.prefix}/wosac/map_based_metrics": final_metrics.map_based_metrics,
+            f"{self.prefix}/wosac/placement_based_metrics": final_metrics.placement_based_metrics,
+            f"{self.prefix}/wosac/min_ade": final_metrics.min_ade,
+            f"{self.prefix}/wosac/scenario_counter": int(self.scenario_counter),
+        }
+        for k in self.field_names:
+            out_dict[f"{self.prefix}/wosac_likelihood/{k}"] = float(metrics_dict[k])
+
+        return out_dict
+
+    @staticmethod
+    def aggregate_metrics_to_buckets(
+        config: long_metrics_pb2.SimAgentMetricsConfig, # type: ignore
+        metrics: long_metrics_pb2.SimAgentMetrics # type: ignore
+    ) -> long_metrics_pb2.SimAgentsBucketedMetrics: # type: ignore
+        """Aggregates metrics into buckets for better readability."""
+        bucketed_metrics = {}
+        for bucket_name, fields_in_bucket in _METRIC_FIELD_NAMES_BY_BUCKET.items():
+            weighted_metric, weights_sum = 0.0, 0.0
+            for field_name in fields_in_bucket:
+                likelihood_field_name = field_name + '_likelihood'
+                weight = getattr(config, field_name).metametric_weight
+                metric_score = getattr(metrics, likelihood_field_name)
+                weighted_metric += weight * metric_score
+                weights_sum += weight
+            if weights_sum == 0:
+                weights_sum = 1  # FIXME: hack!!!
+                # raise ValueError('The bucket\'s weight sum is zero. Check your metrics'
+                #                 ' config.')
+            bucketed_metrics[bucket_name] = weighted_metric / weights_sum
+
+        return long_metrics_pb2.SimAgentsBucketedMetrics(
+                    realism_meta_metric=metrics.metametric,
+                    kinematic_metrics=bucketed_metrics['kinematic'],
+                    interactive_metrics=bucketed_metrics['interactive'],
+                    map_based_metrics=bucketed_metrics['map_based'],
+                    placement_based_metrics=bucketed_metrics['placement_based'],
+                    min_ade=metrics.min_average_displacement_error,
+                    simulated_collision_rate=metrics.simulated_collision_rate,
+                    simulated_offroad_rate=metrics.simulated_offroad_rate,
+                )
+
+    @staticmethod
+    def load_metrics_config(config_path: str = 'dev/metrics/metric_config.textproto',
+                            ) -> long_metrics_pb2.SimAgentMetricsConfig: # type: ignore
+        config = long_metrics_pb2.SimAgentMetricsConfig()
+        with open(config_path, 'r') as f:
+            text_format.Parse(f.read(), config)
+        return config
+
+    def dumps(self, dir):
+        from datetime import datetime
+
+        timestamp = datetime.now().strftime("%m_%d_%H%M%S")
+
+        results = self.compute()
+        path = os.path.join(dir, f'{self.prefix}_{timestamp}.json')
+        with open(path, 'w', encoding='utf-8') as f:
+            json.dump(results, f, indent=4)
+
+        CONSOLE.log(f'Saved results to [bold][yellow]{path}')
+
+
+""" Preprocess Methods """
+
+def _dump_log_metric_features(
+        pkl_dir: str,
+        tfrecords_dir: str,
+        save_dir: str,
+        transform: WaymoTargetBuilder,
+        token_processor: TokenProcessor,
+        scenario_id: str,
+    ):
+
+    try:
+
+        tqdm.write(f'Processing scenario {scenario_id}')
+        save_path = os.path.join(save_dir, f'{scenario_id}.pkl')
+        if os.path.exists(save_path):
+            return
+
+        # load gt data
+        pkl_file = os.path.join(pkl_dir, f'{scenario_id}.pkl')
+        if not os.path.exists(pkl_file):
+            raise FileNotFoundError(f"Not found file {pkl_file}")
+        tfrecord_file = os.path.join(tfrecords_dir, f'{scenario_id}.tfrecords')
+        if not os.path.exists(tfrecord_file):
+            raise FileNotFoundError(f"Not found file {tfrecord_file}")
+
+        scenario_log = scenario_pb2.Scenario()
+        for data in tf.data.TFRecordDataset([tfrecord_file], compression_type=''):
+            scenario_log.ParseFromString(bytes(data.numpy()))
+            break
+
+        with open(pkl_file, 'rb') as f:
+            log_data = pickle.load(f)
+
+        # preprocess data
+        log_data = transform._score_trained_agents(log_data)  # get `train_mask`
+        log_data = token_processor._tokenize_agent(log_data)
+
+        # convert to `JointScene` and compute features
+        log_trajectories = scenario_to_trajectories(scenario_log, processed_scenario=log_data)
+        # log_trajectories = ObjectTrajectories.init_from_processed_scenario(data)
+
+        # NOTE: we do not consider the `evaluation_agent_ids` here
+        # evaluate_agent_ids = torch.tensor(
+        #     get_evaluation_agent_ids(scenario_log)
+        # )
+        evaluate_agent_ids = None
+        log_features = compute_metric_features(
+            log_trajectories, evaluate_agent_ids=evaluate_agent_ids, #scenario_log=scenario_log,
+        )
+
+        # save to pkl file
+        with open(save_path, 'wb') as f:
+            pickle.dump(log_features, f)
+
+    except Exception as e:
+        CONSOLE.log(f'[on red] Failed to process scenario {scenario_id} due to {e}.[/]')
+        return
+
+
+def dump_log_metric_features(log_dir, save_dir):
+
+    buffer_size = 128
+
+    # file loaders
+    pkl_dir = os.path.join(log_dir, 'validation')
+    if not os.path.exists(pkl_dir):
+        raise RuntimeError(f'Not found folder {pkl_dir}')
+    tfrecords_dir = os.path.join(log_dir, 'validation_tfrecords_splitted')
+    if not os.path.exists(tfrecords_dir):
+        raise RuntimeError(f'Not found folder {tfrecords_dir}')
+
+    files = list(fnmatch.filter(os.listdir(pkl_dir), '*.pkl'))
+    json_path = os.path.join(log_dir, 'meta_infos.json')
+    meta_infos = json.load(open(json_path, 'r', encoding='utf-8'))['validation']
+    CONSOLE.log(f"Loaded meta infos from {json_path}")
+    available_scenarios = list(meta_infos.keys())
+    df = pd.DataFrame.from_dict(meta_infos, orient='index')
+    available_scenarios_set = set(available_scenarios)
+    df_filtered = df[(df.index.isin(available_scenarios_set)) & (df['num_agents'] >= 8) & (df['num_agents'] < buffer_size)]
+    valid_scenarios = set(df_filtered.index)
+    files = list(tqdm(filter(lambda fn: fn.removesuffix('.pkl') in valid_scenarios, files), leave=False))
+
+    scenario_ids = list(map(lambda fn: fn.removesuffix('.pkl'), files))
+    CONSOLE.log(f'Loaded {len(scenario_ids)} scenarios from validation split.')
+
+    # initialize
+    transform = WaymoTargetBuilder(num_historical_steps=11,
+                                   num_future_steps=80,
+                                   max_num=32)
+
+    token_processor = TokenProcessor(token_size=2048,
+                                     state_token={'invalid': 0, 'valid': 1, 'enter': 2, 'exit': 3},
+                                     pl2seed_radius=75)
+
+    partial_dump_gt_metric_features = partial(
+        _dump_log_metric_features, pkl_dir, tfrecords_dir, save_dir, transform, token_processor)
+
+    for scenario_id in tqdm(scenario_ids, leave=False, desc='scenarios ...'):
+
+        partial_dump_gt_metric_features(scenario_id)
+
+
+def batch_dump_log_metric_features(log_dir, save_dir, num_workers=64):
+
+    buffer_size = 128
+
+    # file loaders
+    pkl_dir = os.path.join(log_dir, 'validation')
+    if not os.path.exists(pkl_dir):
+        raise RuntimeError(f'Not found folder {pkl_dir}')
+    tfrecords_dir = os.path.join(log_dir, 'validation_tfrecords_splitted')
+    if not os.path.exists(tfrecords_dir):
+        raise RuntimeError(f'Not found folder {tfrecords_dir}')
+
+    files = list(fnmatch.filter(os.listdir(pkl_dir), '*.pkl'))
+    json_path = os.path.join(log_dir, 'meta_infos.json')
+    meta_infos = json.load(open(json_path, 'r', encoding='utf-8'))['validation']
+    CONSOLE.log(f"Loaded meta infos from {json_path}")
+    available_scenarios = list(meta_infos.keys())
+    df = pd.DataFrame.from_dict(meta_infos, orient='index')
+    available_scenarios_set = set(available_scenarios)
+    df_filtered = df[(df.index.isin(available_scenarios_set)) & (df['num_agents'] >= 8) & (df['num_agents'] < buffer_size)]
+    valid_scenarios = set(df_filtered.index)
+    files = list(tqdm(filter(lambda fn: fn.removesuffix('.pkl') in valid_scenarios, files), leave=False))
+
+    scenario_ids = list(map(lambda fn: fn.removesuffix('.pkl'), files))
+    CONSOLE.log(f'Loaded {len(scenario_ids)} scenarios from validation split.')
+
+    # initialize
+    transform = WaymoTargetBuilder(num_historical_steps=11,
+                                   num_future_steps=80,
+                                   max_num=32)
+
+    token_processor = TokenProcessor(token_size=2048,
+                                     state_token={'invalid': 0, 'valid': 1, 'enter': 2, 'exit': 3},
+                                     pl2seed_radius=75)
+
+    partial_dump_gt_metric_features = partial(
+        _dump_log_metric_features, pkl_dir, tfrecords_dir, save_dir, transform, token_processor)
+
+    with multiprocessing.Pool(num_workers) as p:
+        list(tqdm(p.imap_unordered(partial_dump_gt_metric_features, scenario_ids), total=len(scenario_ids)))
+
+
+def aggregate_log_metric_features(load_dir):
+
+    files = list(fnmatch.filter(os.listdir(load_dir), '*.pkl'))
+    if 'total_features.pkl' in files:
+        files.remove('total_features.pkl')
+    CONSOLE.log(f'Loaded {len(files)} scenarios from dumpped log metric features')
+
+    features_fields = [field.name for field in dataclasses.fields(MetricFeatures)]
+    features_fields.remove('object_id')
+
+    # load and append
+    total_features = collections.defaultdict(list)
+    for file in tqdm(files, leave=False, desc='scenario ...'):
+
+        with open(os.path.join(load_dir, file), 'rb') as f:
+            log_features = pickle.load(f)
+
+        for field in features_fields:
+            total_features[field].append(getattr(log_features, field))
+
+    # aggregate
+    features_info = dict()
+    for field in (pbar := tqdm(features_fields, leave=False)):
+        pbar.set_postfix(f=field)
+        if total_features[field][0] is not None:
+            total_features[field] = torch.concat(total_features[field], dim=0)  # n_agent or n_scenario 
+            features_info[field] = total_features[field].shape
+    CONSOLE.log(f'Aggregated log features:\n{features_info}')
+
+    # save
+    save_path = os.path.join(load_dir, 'total_features.pkl')
+    with open(save_path, 'wb') as f:
+        pickle.dump(total_features, f)
+    CONSOLE.log(f'Saved total features to [green]{save_path}.[/]')
+
+    return MetricFeatures(**total_features, object_id=None)
+
+
+def _compute_metrics(
+    metric: LongMetric,
+    load_dir: str,
+    verbose: bool,
+    rollouts_file: str,
+) -> List[long_metrics_pb2.SimAgentMetrics]: # type: ignore
+
+    if verbose:
+        print(f'Processing {rollouts_file}')
+
+    with open(os.path.join(load_dir, rollouts_file), 'rb') as f:
+        rollouts = pickle.load(f)
+    # CONSOLE.log(f'Loaded rollouts from {rollouts_file}')
+
+    return metric.compute_metrics(rollouts)
+
+
+def compute_metrics(load_dir, rollouts_files):
+
+    log_every_n_steps = 100
+
+    metric = LongMetric('val_close_long')
+    CONSOLE.log(f'metrics config:\n{metric.metrics_config}')
+
+    i = 0
+    for rollouts_file in tqdm(rollouts_files, leave=False, desc='Rollouts files ...'):
+
+        # ! compute metrics and update
+        metric.update(
+            metrics=_compute_metrics(metric, load_dir, verbose=False, rollouts_file=rollouts_file)
+        )
+
+        if i % log_every_n_steps == 0:
+            CONSOLE.log(f'Step={i}:\n{metric.compute()}')
+
+        i += 1
+
+    CONSOLE.log(f'[bold][yellow] Compute metrics completed!')
+    CONSOLE.log(f'[bold][yellow] Final metrics: [/]\n {metric.compute()}')
+
+
+def batch_compute_metrics(load_dir, rollouts_files, num_workers, save_dir=None):
+    from queue import Queue
+    from threading import Thread
+
+    if save_dir is None:
+        save_dir = load_dir
+
+    results_buffer = Queue()
+
+    log_every_n_steps = 20
+
+    metric = LongMetric('val_close_long')
+    CONSOLE.log(f'metrics config:\n{metric.metrics_config}')
+
+    def _collect_result():
+        while True:
+            r = results_buffer.get()
+            if r is None:
+                break
+            metric.update(metrics=r)
+            results_buffer.task_done()
+
+    collector = Thread(target=_collect_result, daemon=True)
+    collector.start()
+
+    partial_compute_metrics = partial(_compute_metrics, metric, load_dir, True)
+
+    # ! compute metrics in batch
+    with concurrent.futures.ProcessPoolExecutor(max_workers=num_workers) as executor:
+        # results = list(executor.map(partial_compute_metrics, rollouts_files))
+        futures = [executor.submit(partial_compute_metrics, rollouts_file) for rollouts_file in rollouts_files]
+        # results = [f.result() for f in concurrent.futures.as_completed(futures)]
+
+        for i, future in tqdm(enumerate(concurrent.futures.as_completed(futures)), total=len(futures), leave=False):
+            results_buffer.put(future.result())
+
+            if i % log_every_n_steps == 0:
+                CONSOLE.log(f'Step={i}:\n{metric.compute()}')
+                metric.dumps(save_dir)
+
+    results_buffer.put(None)
+    collector.join()
+
+    CONSOLE.log(f'[bold][yellow] Compute metrics completed!')
+    CONSOLE.log(f'[bold][yellow] Final metrics: [/]\n {metric.compute()}')
+
+    # save results to disk
+    metric.dumps(save_dir)
+
+
+if __name__ == "__main__":
+    parser = ArgumentParser()
+    parser.add_argument('--dump_log', action='store_true')
+    parser.add_argument('--dump_sim', action='store_true')
+    parser.add_argument('--aggregate_log', action='store_true')
+    parser.add_argument('--num_workers', type=int, default=32)
+    parser.add_argument('--compute_metric', action='store_true')
+    parser.add_argument('--log_dir', type=str, default='data/waymo_processed/')
+    parser.add_argument('--sim_dir', type=str, default=None, required=False)
+    parser.add_argument('--save_dir', type=str, default='results', required=False)
+    parser.add_argument('--no_batch', action='store_true')
+    parser.add_argument('--debug', action='store_true')
+    parser.add_argument('--debug_batch', action='store_true')
+    args = parser.parse_args()
+
+    if args.dump_log:
+
+        save_dir = os.path.join(args.log_dir, 'log_features')
+        os.makedirs(save_dir, exist_ok=True)
+
+        if args.no_batch or args.debug:
+            dump_log_metric_features(args.log_dir, save_dir)
+        else:
+            batch_dump_log_metric_features(args.log_dir, save_dir)
+
+    elif args.aggregate_log:
+
+        load_dir = os.path.join(args.log_dir, 'log_features')
+        aggregate_log_metric_features(load_dir)
+
+    elif args.compute_metric:
+
+        assert args.sim_dir is not None and os.path.exists(args.sim_dir), \
+                f'Folder {args.sim_dir} does not exist!'
+        rollouts_files = list(sorted(fnmatch.filter(os.listdir(args.sim_dir), 'idx_*_rollouts.pkl')))
+        CONSOLE.log(f'Found {len(rollouts_files)} rollouts files.')
+
+        os.makedirs(args.save_dir, exist_ok=True)
+        if args.no_batch:
+            compute_metrics(args.sim_dir, rollouts_files)
+
+        else:
+            multiprocessing.set_start_method('spawn', force=True)
+            batch_compute_metrics(args.sim_dir, rollouts_files, args.num_workers, save_dir=args.save_dir)
+
+    elif args.debug:
+
+        debug_path = 'output/scalable_smart_long/validation_catk/idx_0_0_rollouts.pkl'
+
+        # ! for debugging
+        with open(debug_path, 'rb') as f:
+            rollouts = pickle.load(f)
+        metric = LongMetric('debug')
+        CONSOLE.log(f'metrics config: {metric.metrics_config}')
+
+        metric.update(outputs=rollouts)
+        CONSOLE.log(f'metrics:\n{metric.compute()}')
+
+
+    elif args.debug_batch:
+
+        rollouts_files = ['idx_0_rollouts.pkl'] * 1000
+        CONSOLE.log(f'Found {len(rollouts_files)} rollouts files.')
+
+        sim_dir = 'dev/metrics/'
+
+        os.makedirs(args.save_dir, exist_ok=True)
+        multiprocessing.set_start_method('spawn', force=True)
+        batch_compute_metrics(args.sim_dir, rollouts_files, args.num_workers, save_dir=args.save_dir)

seq_10_150_3_3_encode_occ_separate_offsets_bs8_128_no_seqindex_long_lastvalid_ablation_grid/backups/dev/metrics/geometry_utils.py

@@ -0,0 +1,137 @@
+import torch
+import numpy as np
+from torch import Tensor
+from typing import Tuple
+
+
+NUM_VERTICES_IN_BOX = 4
+
+
+def minkowski_sum_of_box_and_box_points(box1_points: Tensor, 
+                                        box2_points: Tensor) -> Tensor:
+    """Batched Minkowski sum of two boxes (counter-clockwise corners in xy)."""
+    point_order_1 = torch.tensor([0, 0, 1, 1, 2, 2, 3, 3], dtype=torch.long)
+    point_order_2 = torch.tensor([0, 1, 1, 2, 2, 3, 3, 0], dtype=torch.long)
+
+    box1_start_idx, downmost_box1_edge_direction = _get_downmost_edge_in_box(
+        box1_points)
+    box2_start_idx, downmost_box2_edge_direction = _get_downmost_edge_in_box(
+        box2_points)
+
+    condition = (cross_product_2d(downmost_box1_edge_direction, downmost_box2_edge_direction) >= 0.)
+    condition = condition.repeat(1, 8)
+
+    box1_point_order = torch.where(condition, point_order_2, point_order_1)
+    box1_point_order = (box1_point_order + box1_start_idx) % NUM_VERTICES_IN_BOX
+    ordered_box1_points = torch.gather(
+        box1_points, 1, box1_point_order.unsqueeze(-1).expand(-1, -1, 2))
+
+    box2_point_order = torch.where(condition, point_order_1, point_order_2)
+    box2_point_order = (box2_point_order + box2_start_idx) % NUM_VERTICES_IN_BOX
+    ordered_box2_points = torch.gather(
+        box2_points, 1, box2_point_order.unsqueeze(-1).expand(-1, -1, 2))
+
+    minkowski_sum = ordered_box1_points + ordered_box2_points
+
+    return minkowski_sum
+
+
+def signed_distance_from_point_to_convex_polygon(query_points: Tensor, polygon_points: Tensor) -> Tensor:
+    """Finds the signed distances from query points to convex polygons."""
+    tangent_unit_vectors, normal_unit_vectors, edge_lengths = _get_edge_info(
+        polygon_points)
+
+    query_points = query_points.unsqueeze(1)
+    vertices_to_query_vectors = query_points - polygon_points
+    vertices_distances = torch.norm(vertices_to_query_vectors, dim=-1)
+
+    edge_signed_perp_distances = torch.sum(-normal_unit_vectors * vertices_to_query_vectors, dim=-1)
+
+    is_inside = torch.all(edge_signed_perp_distances <= 0, dim=-1)
+
+    projection_along_tangent = torch.sum(tangent_unit_vectors * vertices_to_query_vectors, dim=-1)
+    projection_along_tangent_proportion = projection_along_tangent / edge_lengths
+
+    is_projection_on_edge = (projection_along_tangent_proportion >= 0.) & (
+                             projection_along_tangent_proportion <= 1.)
+
+    edge_perp_distances = edge_signed_perp_distances.abs()
+    edge_distances = torch.where(is_projection_on_edge, edge_perp_distances, torch.tensor(np.inf))
+
+    edge_and_vertex_distance = torch.cat([edge_distances, vertices_distances], dim=-1)
+    min_distance = torch.min(edge_and_vertex_distance, dim=-1)[0]
+
+    signed_distances = torch.where(is_inside, -min_distance, min_distance)
+
+    return signed_distances
+
+
+def _get_downmost_edge_in_box(box: Tensor) -> Tuple[Tensor, Tensor]:
+    """Finds the downmost (lowest y-coordinate) edge in the box."""
+    downmost_vertex_idx = torch.argmin(box[..., 1], dim=-1, keepdim=True)
+
+    edge_start_vertex = torch.gather(box, 1, downmost_vertex_idx.unsqueeze(-1).expand(-1, -1, 2))
+    edge_end_idx = (downmost_vertex_idx + 1) % NUM_VERTICES_IN_BOX
+    edge_end_vertex = torch.gather(box, 1, edge_end_idx.unsqueeze(-1).expand(-1, -1, 2))
+    
+    downmost_edge = edge_end_vertex - edge_start_vertex
+    downmost_edge_length = torch.norm(downmost_edge, dim=-1, keepdim=True)
+    downmost_edge_direction = downmost_edge / downmost_edge_length
+    
+    return downmost_vertex_idx, downmost_edge_direction
+
+
+def cross_product_2d(a: Tensor, b: Tensor) -> Tensor:
+    return a[..., 0] * b[..., 1] - a[..., 1] * b[..., 0]
+
+
+def dot_product_2d(a: Tensor, b: Tensor) -> Tensor:
+    return a[..., 0] * b[..., 0] + a[..., 1] * b[..., 1]
+
+
+def _get_edge_info(polygon_points: Tensor):
+    """
+    Computes properties about the edges of a polygon.
+
+    Args:
+        polygon_points: Tensor containing the vertices of each polygon, with
+          shape (num_polygons, num_points_per_polygon, 2). Each polygon is assumed
+          to have an equal number of vertices.
+
+    Returns:
+        tangent_unit_vectors: A unit vector in (x,y) with the same direction as
+          the tangent to the edge. Shape: (num_polygons, num_points_per_polygon, 2).
+        normal_unit_vectors: A unit vector in (x,y) with the same direction as
+          the normal to the edge.
+          Shape: (num_polygons, num_points_per_polygon, 2).
+        edge_lengths: Lengths of the edges.
+          Shape (num_polygons, num_points_per_polygon).
+    """
+    # Shift the polygon points by 1 position to get the edges.
+    first_point_in_polygon = polygon_points[:, 0:1, :]  # Shape: (num_polygons, 1, 2)
+    shifted_polygon_points = torch.cat([polygon_points[:, 1:, :], first_point_in_polygon], dim=1)  
+    # Shape: (num_polygons, num_points_per_polygon, 2)
+
+    edge_vectors = shifted_polygon_points - polygon_points  # Shape: (num_polygons, num_points_per_polygon, 2)
+    edge_lengths = torch.norm(edge_vectors, dim=-1)  # Shape: (num_polygons, num_points_per_polygon)
+
+    # Avoid division by zero by adding a small epsilon
+    eps = torch.finfo(edge_lengths.dtype).eps
+    tangent_unit_vectors = edge_vectors / (edge_lengths[..., None] + eps)  # Shape: (num_polygons, num_points_per_polygon, 2)
+
+    normal_unit_vectors = torch.stack(
+        [-tangent_unit_vectors[..., 1], tangent_unit_vectors[..., 0]], dim=-1
+    )  # Shape: (num_polygons, num_points_per_polygon, 2)
+
+    return tangent_unit_vectors, normal_unit_vectors, edge_lengths
+
+
+def rotate_2d_points(xys: Tensor, rotation_yaws: Tensor) -> Tensor:
+    """Rotates `xys` counterclockwise using the `rotation_yaws`."""
+    cos_yaws = torch.cos(rotation_yaws)
+    sin_yaws = torch.sin(rotation_yaws)
+    
+    rotated_x = cos_yaws * xys[..., 0] - sin_yaws * xys[..., 1]
+    rotated_y = sin_yaws * xys[..., 0] + cos_yaws * xys[..., 1]
+    
+    return torch.stack([rotated_x, rotated_y], axis=-1)

seq_10_150_3_3_encode_occ_separate_offsets_bs8_128_no_seqindex_long_lastvalid_ablation_grid/backups/dev/metrics/interact_features.py

@@ -0,0 +1,220 @@
+import math
+import torch
+from torch import Tensor
+
+from dev.metrics import box_utils
+from dev.metrics import geometry_utils
+from dev.metrics import trajectory_features
+
+
+EXTREMELY_LARGE_DISTANCE = 1e10
+COLLISION_DISTANCE_THRESHOLD = 0.0
+CORNER_ROUNDING_FACTOR = 0.7
+MAX_HEADING_DIFF = math.radians(75.0)
+MAX_HEADING_DIFF_FOR_SMALL_OVERLAP = math.radians(10.0)
+SMALL_OVERLAP_THRESHOLD = 0.5
+MAXIMUM_TIME_TO_COLLISION = 5.0
+
+
+def compute_distance_to_nearest_object(
+    center_x: Tensor,
+    center_y: Tensor,
+    center_z: Tensor,
+    length: Tensor,
+    width: Tensor,
+    height: Tensor,
+    heading: Tensor,
+    valid: Tensor,
+    evaluated_object_mask: Tensor,
+    corner_rounding_factor: float = CORNER_ROUNDING_FACTOR,
+) -> Tensor:
+    """Computes the distance to nearest object for each of the evaluated objects."""
+    boxes = torch.stack([center_x, center_y, center_z, length, width, height, heading], dim=-1)
+    num_objects, num_steps, num_features = boxes.shape
+
+    shrinking_distance = (torch.minimum(boxes[:, :, 3], boxes[:, :, 4]) * corner_rounding_factor / 2.)
+
+    boxes = torch.cat([
+        boxes[:, :, :3],
+        boxes[:, :, 3:4] - 2.0 * shrinking_distance[..., None],
+        boxes[:, :, 4:5] - 2.0 * shrinking_distance[..., None],
+        boxes[:, :, 5:]
+    ], dim=2)
+
+    boxes = boxes.reshape(num_objects * num_steps, num_features)
+
+    box_corners = box_utils.get_upright_3d_box_corners(boxes)[:, :4, :2]
+    box_corners = box_corners.reshape(num_objects, num_steps, 4, 2)
+
+    eval_corners = box_corners[evaluated_object_mask]
+    num_eval_objects = eval_corners.shape[0]
+    other_corners = box_corners[~evaluated_object_mask]
+    all_corners = torch.cat([eval_corners, other_corners], dim=0)
+
+    eval_corners = eval_corners.unsqueeze(1).expand(num_eval_objects, num_objects, num_steps, 4, 2)
+    all_corners = all_corners.unsqueeze(0).expand(num_eval_objects, num_objects, num_steps, 4, 2)
+
+    eval_corners = eval_corners.reshape(num_eval_objects * num_objects * num_steps, 4, 2)
+    all_corners = all_corners.reshape(num_eval_objects * num_objects * num_steps, 4, 2)
+
+    neg_all_corners = -1.0 * all_corners
+    minkowski_sum = geometry_utils.minkowski_sum_of_box_and_box_points(
+        box1_points=eval_corners, box2_points=neg_all_corners,
+    )
+
+    assert minkowski_sum.shape[1:] == (8, 2), f"Shape mismatch: {minkowski_sum.shape}, expected (*, 8, 2)"
+    signed_distances_flat = (
+        geometry_utils.signed_distance_from_point_to_convex_polygon(
+            query_points=torch.zeros_like(minkowski_sum[:, 0, :]),
+            polygon_points=minkowski_sum,
+        )
+    )
+
+    signed_distances = signed_distances_flat.reshape(num_eval_objects, num_objects, num_steps)
+
+    eval_shrinking_distance = shrinking_distance[evaluated_object_mask]
+    other_shrinking_distance = shrinking_distance[~evaluated_object_mask]
+    all_shrinking_distance = torch.cat([eval_shrinking_distance, other_shrinking_distance], dim=0)
+
+    signed_distances -= eval_shrinking_distance.unsqueeze(1)
+    signed_distances -= all_shrinking_distance.unsqueeze(0)
+
+    self_mask = torch.eye(num_eval_objects, num_objects, dtype=torch.float32)[:, :, None]
+    signed_distances = signed_distances + self_mask * EXTREMELY_LARGE_DISTANCE
+
+    eval_validity = valid[evaluated_object_mask]
+    other_validity = valid[~evaluated_object_mask]
+    all_validity = torch.cat([eval_validity, other_validity], dim=0)
+
+    valid_mask = eval_validity.unsqueeze(1) & all_validity.unsqueeze(0)
+
+    signed_distances = torch.where(valid_mask, signed_distances, EXTREMELY_LARGE_DISTANCE)
+
+    return torch.min(signed_distances, dim=1).values
+
+
+def compute_time_to_collision_with_object_in_front(
+    *,
+    center_x: Tensor,
+    center_y: Tensor,
+    length: Tensor,
+    width: Tensor,
+    heading: Tensor,
+    valid: Tensor,
+    evaluated_object_mask: Tensor,
+    seconds_per_step: float,
+) -> Tensor:
+    """Computes the time-to-collision of the evaluated objects."""
+    # `speed` shape: (num_objects, num_steps)
+    speed = trajectory_features.compute_kinematic_features(
+        x=center_x,
+        y=center_y,
+        z=torch.zeros_like(center_x),
+        heading=heading,
+        seconds_per_step=seconds_per_step,
+    )[0]
+
+    boxes = torch.stack([center_x, center_y, length, width, heading, speed], dim=-1)
+    boxes = boxes.permute(1, 0, 2)  # (num_steps, num_objects, 6)
+    valid = valid.permute(1, 0)
+
+    eval_boxes = boxes[:, evaluated_object_mask]
+    ego_xy, ego_sizes, ego_yaw, ego_speed = torch.split(eval_boxes, [2, 2, 1, 1], dim=-1)
+    other_xy, other_sizes, other_yaw, _ = torch.split(boxes, [2, 2, 1, 1], dim=-1)
+
+    yaw_diff = torch.abs(other_yaw[:, None] - ego_yaw[:, :, None])
+    yaw_diff_cos = torch.cos(yaw_diff)
+    yaw_diff_sin = torch.sin(yaw_diff)
+
+    other_long_offset = geometry_utils.dot_product_2d(
+        other_sizes[:, None] / 2.0, torch.abs(torch.cat([yaw_diff_cos, yaw_diff_sin], dim=-1))
+    )
+    other_lat_offset = geometry_utils.dot_product_2d(
+        other_sizes[:, None] / 2.0, torch.abs(torch.cat([yaw_diff_sin, yaw_diff_cos], dim=-1))
+    )
+
+    other_relative_xy = geometry_utils.rotate_2d_points(
+        (other_xy[:, None] - ego_xy[:, :, None]), -ego_yaw
+    )
+    
+    long_distance = (
+        other_relative_xy[..., 0] - ego_sizes[:, :, None, 0] / 2.0 - other_long_offset
+    )
+    lat_overlap = (
+        torch.abs(other_relative_xy[..., 1]) - ego_sizes[:, :, None, 1] / 2.0 - other_lat_offset
+    )
+    
+    following_mask = _get_object_following_mask(
+        long_distance, lat_overlap, yaw_diff[..., 0]
+    )
+    valid_mask = valid[:, None] & following_mask
+    
+    masked_long_distance = (
+        long_distance + (1.0 - valid_mask.float()) * EXTREMELY_LARGE_DISTANCE
+    )
+    
+    box_ahead_index = masked_long_distance.argmin(dim=-1)
+    distance_to_box_ahead = torch.gather(
+        masked_long_distance, -1, box_ahead_index.unsqueeze(-1)
+    ).squeeze(-1)
+    
+    speed_broadcast = speed.T[:, None, :].expand_as(masked_long_distance)
+    box_ahead_speed = torch.gather(speed_broadcast, -1, box_ahead_index.unsqueeze(-1)).squeeze(-1)
+    
+    rel_speed = ego_speed[..., 0] - box_ahead_speed
+    time_to_collision = torch.where(
+        rel_speed > 0.0,
+        torch.minimum(distance_to_box_ahead / rel_speed,
+                      torch.tensor(MAXIMUM_TIME_TO_COLLISION)),  # the float will be broadcasted automatically
+        MAXIMUM_TIME_TO_COLLISION,
+    )
+    
+    return time_to_collision.T
+
+
+def _get_object_following_mask(
+    longitudinal_distance: Tensor,
+    lateral_overlap: Tensor,
+    yaw_diff: Tensor,
+) -> Tensor:
+    """Checks whether objects satisfy criteria for following another object.
+
+    An object on which the criteria are applied is called "ego object" in this
+    function to disambiguate it from the other objects acting as obstacles.
+
+    An "ego" object is considered to be following another object if they satisfy
+    conditions on the longitudinal distance, lateral overlap, and yaw alignment
+    between them.
+
+    Args:
+        longitudinal_distance: A float Tensor with shape (batch_dim, num_egos,
+          num_others) containing longitudinal distances from the back side of each
+          ego box to every other box.
+        lateral_overlap: A float Tensor with shape (batch_dim, num_egos, num_others)
+          containing lateral overlaps of other boxes over the trails of ego boxes.
+        yaw_diff: A float Tensor with shape (batch_dim, num_egos, num_others)
+          containing absolute yaw differences between egos and other boxes.
+
+    Returns:
+        A boolean Tensor with shape (batch_dim, num_egos, num_others) indicating for
+        each ego box if it is following the other boxes.
+    """
+    # Check object is ahead of the ego box's front.
+    valid_mask = longitudinal_distance > 0.0
+
+    # Check alignment.
+    valid_mask = torch.logical_and(valid_mask, yaw_diff <= MAX_HEADING_DIFF)
+
+    # Check object is directly ahead of the ego box.
+    valid_mask = torch.logical_and(valid_mask, lateral_overlap < 0.0)
+
+    # Check strict alignment if the overlap is small.
+    # `lateral_overlap` is a signed penetration distance: it is negative when the
+    # boxes have an actual lateral overlap.
+    return torch.logical_and(
+        valid_mask,
+        torch.logical_or(
+            lateral_overlap < -SMALL_OVERLAP_THRESHOLD,
+            yaw_diff <= MAX_HEADING_DIFF_FOR_SMALL_OVERLAP,
+        ),
+    )

seq_10_150_3_3_encode_occ_separate_offsets_bs8_128_no_seqindex_long_lastvalid_ablation_grid/backups/dev/metrics/map_features.py

@@ -0,0 +1,349 @@
+import torch
+from torch import Tensor
+from typing import Optional, Sequence
+
+from dev.metrics import box_utils
+from dev.metrics import geometry_utils
+from dev.metrics.protos import map_pb2
+
+# Constant distance to apply when distances are invalid. This will avoid the
+# propagation of nans and should be reduced out when taking the minimum anyway.
+EXTREMELY_LARGE_DISTANCE = 1e10
+# Off-road threshold, i.e. smallest distance away from the road edge that is
+# considered to be a off-road.
+OFFROAD_DISTANCE_THRESHOLD = 0.0
+
+# How close the start and end point of a map feature need to be for the feature
+# to be considered cyclic, in m^2.
+_CYCLIC_MAP_FEATURE_TOLERANCE_M2 = 1.0
+# Scaling factor for vertical distances used when finding the closest segment to
+# a query point. This prevents wrong associations in cases with under- and
+# over-passes.
+_Z_STRETCH_FACTOR = 3.0
+
+_Polyline = Sequence[map_pb2.MapPoint]
+
+
+def compute_distance_to_road_edge(
+    *,
+    center_x: Tensor,
+    center_y: Tensor,
+    center_z: Tensor,
+    length: Tensor,
+    width: Tensor,
+    height: Tensor,
+    heading: Tensor,
+    valid: Tensor,
+    evaluated_object_mask: Tensor,
+    road_edge_polylines: Sequence[_Polyline],
+) -> Tensor:
+    """Computes the distance to the road edge for each of the evaluated objects."""
+    if not road_edge_polylines:
+        raise ValueError('Missing road edges.')
+
+    # Concatenate tensors to have the same convention as `box_utils`.
+    boxes = torch.stack([center_x, center_y, center_z, length, width, height, heading], dim=-1)
+    num_objects, num_steps, num_features = boxes.shape
+    boxes = boxes.reshape(num_objects * num_steps, num_features)
+
+    # Compute box corners using `box_utils`, and take the xyz coords of the bottom corners.
+    box_corners = box_utils.get_upright_3d_box_corners(boxes)[:, :4]
+    box_corners = box_corners.reshape(num_objects, num_steps, 4, 3)
+
+    # Gather objects in the evaluation set
+    eval_corners = box_corners[evaluated_object_mask]
+    num_eval_objects = eval_corners.shape[0]
+
+    # Flatten query points.
+    flat_eval_corners = eval_corners.reshape(-1, 3)
+
+    # Tensorize road edges.
+    polylines_tensor = _tensorize_polylines(road_edge_polylines)
+    is_polyline_cyclic = _check_polyline_cycles(road_edge_polylines)
+
+    # Compute distances for all query points.
+    corner_distance_to_road_edge = _compute_signed_distance_to_polylines(
+        xyzs=flat_eval_corners, polylines=polylines_tensor,
+        is_polyline_cyclic=is_polyline_cyclic, z_stretch=_Z_STRETCH_FACTOR
+    )
+
+    # Reshape back to (num_evaluated_objects, num_steps, 4)
+    corner_distance_to_road_edge = corner_distance_to_road_edge.reshape(num_eval_objects, num_steps, 4)
+
+    # Reduce to most off-road corner.
+    signed_distances = torch.max(corner_distance_to_road_edge, dim=-1)[0]
+
+    # Mask out invalid boxes.
+    eval_validity = valid[evaluated_object_mask]
+
+    return torch.where(eval_validity, signed_distances, -EXTREMELY_LARGE_DISTANCE)
+
+
+def _tensorize_polylines(polylines):
+    """Stacks a sequence of polylines into a tensor.
+
+    Args:
+        polylines: A sequence of Polyline objects.
+
+    Returns:
+        A float tensor with shape (num_polylines, max_length, 4) containing xyz
+        coordinates and a validity flag for all points in the polylines. Polylines
+        are padded with zeros up to the length of the longest one.
+    """
+    polyline_tensors = []
+    max_length = 0
+
+    for polyline in polylines:
+        # Skip degenerate polylines.
+        if len(polyline) < 2:
+            continue
+        max_length = max(max_length, len(polyline))
+        polyline_tensors.append(
+            torch.tensor(
+                [[map_point.x, map_point.y, map_point.z, 1.0] for map_point in polyline], 
+                dtype=torch.float32
+            )
+        )
+
+    # Pad and stack polylines
+    padded_polylines = [
+        torch.cat([p, torch.zeros((max_length - p.shape[0], 4), dtype=torch.float32)], dim=0)
+        for p in polyline_tensors
+    ]
+
+    return torch.stack(padded_polylines, dim=0)
+
+
+def _check_polyline_cycles(polylines):
+    """Checks if given polylines are cyclic and returns the result as a tensor.
+
+    Args:
+        polylines: A sequence of Polyline objects.
+        tolerance: A float representing the cyclic tolerance.
+
+    Returns:
+        A bool tensor with shape (num_polylines) indicating whether each polyline is cyclic.
+    """
+    cycles = []
+    for polyline in polylines:
+        # Skip degenerate polylines.
+        if len(polyline) < 2:
+            continue
+        first_point = torch.tensor([polyline[0].x, polyline[0].y, polyline[0].z], dtype=torch.float32)
+        last_point = torch.tensor([polyline[-1].x, polyline[-1].y, polyline[-1].z], dtype=torch.float32)
+        cycles.append(torch.sum((first_point - last_point) ** 2) < _CYCLIC_MAP_FEATURE_TOLERANCE_M2)
+
+    return torch.stack(cycles, dim=0)
+
+
+def _compute_signed_distance_to_polylines(
+    xyzs: Tensor,
+    polylines: Tensor,
+    is_polyline_cyclic: Optional[Tensor] = None,
+    z_stretch: float = 1.0,
+) -> Tensor:
+    """Computes the signed distance to the 2D boundary defined by polylines.
+
+    Negative distances correspond to being inside the boundary (e.g. on the
+    road), positive distances to being outside (e.g. off-road).
+
+    The polylines should be oriented such that port side is inside the boundary
+    and starboard is outside, a.k.a counterclockwise winding order.
+
+    The altitudes i.e. the z-coordinates of query points and polyline segments
+    are used to pair each query point with the most relevant segment, that is
+    closest and at the right altitude. The distances returned are 2D distances in
+    the xy plane.
+
+    Args:
+      xyzs: A float Tensor of shape (num_points, 3) containing xyz coordinates of
+        query points.
+      polylines: Tensor with shape (num_polylines, num_segments+1, 4) containing
+        sequences of xyz coordinates and validity, representing start and end
+        points of consecutive segments.
+      is_polyline_cyclic: A boolean Tensor with shape (num_polylines) indicating
+        whether each polyline is cyclic. If None, all polylines are considered
+        non-cyclic.
+      z_stretch: Factor by which to scale distances over the z axis. This can be
+        done to ensure edge points from the wrong level (e.g. overpasses) are not
+        selected. Defaults to 1.0 (no stretching).
+
+    Returns:
+      A tensor of shape (num_points), containing the signed 2D distance from
+      queried points to the nearest polyline.
+    """
+    num_points = xyzs.shape[0]
+    assert xyzs.shape == (num_points, 3), f"Expected shape ({num_points}, 3), but got {xyzs.shape}"
+    num_polylines = polylines.shape[0]
+    num_segments = polylines.shape[1] - 1
+    assert polylines.shape == (num_polylines, num_segments + 1, 4), \
+        f"Expected shape ({num_polylines}, {num_segments + 1}, 4), but got {polylines.shape}"
+
+    # shape: (num_polylines, num_segments+1)
+    is_point_valid = polylines[:, :, 3].bool()
+    # shape: (num_polylines, num_segments)
+    is_segment_valid = is_point_valid[:, :-1] & is_point_valid[:, 1:]
+
+    if is_polyline_cyclic is None:
+        is_polyline_cyclic = torch.zeros(num_polylines, dtype=torch.bool)
+    else:
+        assert is_polyline_cyclic.shape == (num_polylines,), \
+            f"Expected shape ({num_polylines},), but got {is_polyline_cyclic.shape}"
+
+    # Get distance to each segment.
+    # shape: (num_points, num_polylines, num_segments, 3)
+    xyz_starts = polylines[None, :, :-1, :3]
+    xyz_ends = polylines[None, :, 1:, :3]
+    start_to_point = xyzs[:, None, None, :3] - xyz_starts
+    start_to_end = xyz_ends - xyz_starts
+
+    # Relative coordinate of point projection on segment.
+    # shape: (num_points, num_polylines, num_segments)
+    numerator = geometry_utils.dot_product_2d(
+        start_to_point[..., :2], start_to_end[..., :2]
+    )
+    denominator = geometry_utils.dot_product_2d(
+        start_to_end[..., :2], start_to_end[..., :2]
+    )
+    rel_t = torch.where(denominator != 0, numerator / denominator, torch.zeros_like(numerator))
+
+    # Negative if point is on port side of segment, positive if point on
+    # starboard side of segment.
+    # shape: (num_points, num_polylines, num_segments)
+    n = torch.sign(
+        geometry_utils.cross_product_2d(
+            start_to_point[..., :2], start_to_end[..., :2]
+        )
+    )
+
+    # Compute the absolute 3D distance to segment.
+    # The vertical component is scaled by `z-stretch` to increase the separation
+    # between different road altitudes.
+    # shape: (num_points, num_polylines, num_segments, 3)
+    segment_to_point = start_to_point - (
+        start_to_end * torch.clamp(rel_t, 0.0, 1.0)[..., None]
+    )
+    stretch_vector = torch.tensor([1.0, 1.0, z_stretch], dtype=torch.float32)
+    distance_to_segment_3d = torch.norm(
+        segment_to_point * stretch_vector[None, None, None],
+        dim=-1,
+    )
+    
+    # Absolute planar distance to segment.
+    # shape: (num_points, num_polylines, num_segments)
+    distance_to_segment_2d = torch.norm(segment_to_point[..., :2], dim=-1)
+
+    # Padded start-to-end segments.
+    # shape: (num_points, num_polylines, num_segments+2, 2)
+    start_to_end_padded = torch.cat(
+        [
+            start_to_end[:, :, -1:, :2],
+            start_to_end[..., :2],
+            start_to_end[:, :, :1, :2],
+        ],
+        dim=-2,
+    )
+
+    # shape: (num_points, num_polylines, num_segments+1)
+    is_locally_convex = torch.gt(
+        geometry_utils.cross_product_2d(
+            start_to_end_padded[:, :, :-1], start_to_end_padded[:, :, 1:]
+        ),
+        0.,
+    )
+
+    # Get shifted versions of `n` and `is_segment_valid`. If the polyline is
+    # cyclic, the tensors are rolled, else they are padded with their edge value.
+    # shape: (num_points, num_polylines, num_segments)
+    n_prior = torch.cat(
+        [
+            torch.where(
+                is_polyline_cyclic[None, :, None],
+                n[:, :, -1:],
+                n[:, :, :1],
+            ),
+            n[:, :, :-1],
+        ],
+        dim=-1,
+    )
+    n_next = torch.cat(
+        [
+            n[:, :, 1:],
+            torch.where(
+                is_polyline_cyclic[None, :, None],
+                n[:, :, :1],
+                n[:, :, -1:],
+            ),
+        ],
+        dim=-1,
+    )
+    # shape: (num_polylines, num_segments)
+    is_prior_segment_valid = torch.cat(
+        [
+            torch.where(
+                is_polyline_cyclic[:, None],
+                is_segment_valid[:, -1:],
+                is_segment_valid[:, :1],
+            ),
+            is_segment_valid[:, :-1],
+        ],
+        dim=-1,
+    )
+    is_next_segment_valid = torch.cat(
+        [
+            is_segment_valid[:, 1:],
+            torch.where(
+                is_polyline_cyclic[:, None],
+                is_segment_valid[:, :1],
+                is_segment_valid[:, -1:],
+            ),
+        ],
+        dim=-1,
+    )
+
+    # shape: (num_points, num_polylines, num_segments)
+    sign_if_before = torch.where(
+        is_locally_convex[:, :, :-1],
+        torch.maximum(n, n_prior),
+        torch.minimum(n, n_prior),
+    )
+    sign_if_after = torch.where(
+        is_locally_convex[:, :, 1:], torch.maximum(n, n_next), torch.minimum(n, n_next)
+    )
+
+    # shape: (num_points, num_polylines, num_segments)
+    sign_to_segment = torch.where(
+        (rel_t < 0.0) & is_prior_segment_valid,
+        sign_if_before,
+        torch.where((rel_t > 1.0) & is_next_segment_valid, sign_if_after, n),
+    )
+
+    # Flatten polylines together.
+    # shape: (num_points, all_segments)
+    distance_to_segment_3d = distance_to_segment_3d.view(num_points, num_polylines * num_segments)
+    distance_to_segment_2d = distance_to_segment_2d.view(num_points, num_polylines * num_segments)
+    sign_to_segment = sign_to_segment.view(num_points, num_polylines * num_segments)
+
+    # Mask out invalid segments.
+    # shape: (all_segments)
+    is_segment_valid = is_segment_valid.view(num_polylines * num_segments)
+    # shape: (num_points, all_segments)
+    distance_to_segment_3d = torch.where(
+        is_segment_valid[None],
+        distance_to_segment_3d,
+        EXTREMELY_LARGE_DISTANCE,
+    )
+    distance_to_segment_2d = torch.where(
+        is_segment_valid[None],
+        distance_to_segment_2d,
+        EXTREMELY_LARGE_DISTANCE,
+    )
+
+    # Get closest segment according to absolute 3D distance and return the
+    # corresponding signed 2D distance.
+    # shape: (num_points)
+    closest_segment_index = torch.argmin(distance_to_segment_3d, dim=-1)
+    distance_sign = torch.gather(sign_to_segment, 1, closest_segment_index.unsqueeze(-1)).squeeze(-1)
+    distance_2d = torch.gather(distance_to_segment_2d, 1, closest_segment_index.unsqueeze(-1)).squeeze(-1)
+
+    return distance_sign * distance_2d

seq_10_150_3_3_encode_occ_separate_offsets_bs8_128_no_seqindex_long_lastvalid_ablation_grid/backups/dev/metrics/placement_features.py

@@ -0,0 +1,48 @@
+import torch
+from torch import Tensor
+from typing import Optional, Sequence, List
+
+
+def compute_num_placement(
+    valid: Tensor,  # [n_agent, n_step]
+    state: Tensor,  # [n_agent, n_step]
+    av_id: int,
+    object_id: Tensor,
+    agent_state: List[str],
+) -> Tensor:
+
+    enter_state = agent_state.index('enter')
+    exit_state = agent_state.index('exit')
+
+    av_index = object_id.tolist().index(av_id)
+    state[av_index] = -1  # we do not incorporate the sdc
+
+    is_bos = state == enter_state
+    is_eos = state == exit_state
+
+    num_bos = torch.sum(is_bos, dim=0)
+    num_eos = torch.sum(is_eos, dim=0)
+
+    return num_bos, num_eos
+
+
+def compute_distance_placement(
+    position: Tensor,
+    state: Tensor,
+    valid: Tensor,
+    av_id: int,
+    object_id: Tensor,
+    agent_state: List[str],
+) -> Tensor:
+
+    enter_state = agent_state.index('enter')
+    exit_state = agent_state.index('exit')
+
+    av_index = object_id.tolist().index(av_id)
+    state[av_index] = -1  # we do not incorporate the sdc
+    distance = torch.norm(position - position[av_index : av_index + 1], p=2, dim=-1)
+
+    bos_distance = distance * (state == enter_state)
+    eos_distance = distance * (state == exit_state)
+
+    return bos_distance, eos_distance

seq_10_150_3_3_encode_occ_separate_offsets_bs8_128_no_seqindex_long_lastvalid_ablation_grid/backups/dev/metrics/protos/long_metrics_pb2.py

@@ -0,0 +1,648 @@
+# -*- coding: utf-8 -*-
+# Generated by the protocol buffer compiler.  DO NOT EDIT!
+# source: dev/metrics/protos/long_metrics.proto
+
+from google.protobuf import descriptor as _descriptor
+from google.protobuf import message as _message
+from google.protobuf import reflection as _reflection
+from google.protobuf import symbol_database as _symbol_database
+# @@protoc_insertion_point(imports)
+
+_sym_db = _symbol_database.Default()
+
+
+
+
+DESCRIPTOR = _descriptor.FileDescriptor(
+  name='dev/metrics/protos/long_metrics.proto',
+  package='long_metric',
+  syntax='proto2',
+  serialized_options=None,
+  create_key=_descriptor._internal_create_key,
+  serialized_pb=b'\n%dev/metrics/protos/long_metrics.proto\x12\x0blong_metric\"\xb4\x0c\n\x15SimAgentMetricsConfig\x12\x46\n\x0clinear_speed\x18\x01 \x01(\x0b\x32\x30.long_metric.SimAgentMetricsConfig.FeatureConfig\x12M\n\x13linear_acceleration\x18\x02 \x01(\x0b\x32\x30.long_metric.SimAgentMetricsConfig.FeatureConfig\x12G\n\rangular_speed\x18\x03 \x01(\x0b\x32\x30.long_metric.SimAgentMetricsConfig.FeatureConfig\x12N\n\x14\x61ngular_acceleration\x18\x04 \x01(\x0b\x32\x30.long_metric.SimAgentMetricsConfig.FeatureConfig\x12T\n\x1a\x64istance_to_nearest_object\x18\x05 \x01(\x0b\x32\x30.long_metric.SimAgentMetricsConfig.FeatureConfig\x12N\n\x14\x63ollision_indication\x18\x06 \x01(\x0b\x32\x30.long_metric.SimAgentMetricsConfig.FeatureConfig\x12K\n\x11time_to_collision\x18\x07 \x01(\x0b\x32\x30.long_metric.SimAgentMetricsConfig.FeatureConfig\x12O\n\x15\x64istance_to_road_edge\x18\x08 \x01(\x0b\x32\x30.long_metric.SimAgentMetricsConfig.FeatureConfig\x12L\n\x12offroad_indication\x18\t \x01(\x0b\x32\x30.long_metric.SimAgentMetricsConfig.FeatureConfig\x12G\n\rnum_placement\x18\n \x01(\x0b\x32\x30.long_metric.SimAgentMetricsConfig.FeatureConfig\x12H\n\x0enum_removement\x18\x0b \x01(\x0b\x32\x30.long_metric.SimAgentMetricsConfig.FeatureConfig\x12L\n\x12\x64istance_placement\x18\x0c \x01(\x0b\x32\x30.long_metric.SimAgentMetricsConfig.FeatureConfig\x12M\n\x13\x64istance_removement\x18\r \x01(\x0b\x32\x30.long_metric.SimAgentMetricsConfig.FeatureConfig\x1a\xc0\x02\n\rFeatureConfig\x12I\n\thistogram\x18\x01 \x01(\x0b\x32\x34.long_metric.SimAgentMetricsConfig.HistogramEstimateH\x00\x12R\n\x0ekernel_density\x18\x02 \x01(\x0b\x32\x38.long_metric.SimAgentMetricsConfig.KernelDensityEstimateH\x00\x12I\n\tbernoulli\x18\x03 \x01(\x0b\x32\x34.long_metric.SimAgentMetricsConfig.BernoulliEstimateH\x00\x12\x1d\n\x15independent_timesteps\x18\x04 \x01(\x08\x12\x19\n\x11metametric_weight\x18\x05 \x01(\x02\x42\x0b\n\testimator\x1av\n\x11HistogramEstimate\x12\x0f\n\x07min_val\x18\x01 \x01(\x02\x12\x0f\n\x07max_val\x18\x02 \x01(\x02\x12\x10\n\x08num_bins\x18\x03 \x01(\x05\x12-\n\x1e\x61\x64\x64itive_smoothing_pseudocount\x18\x04 \x01(\x02:\x05\x30.001\x1a*\n\x15KernelDensityEstimate\x12\x11\n\tbandwidth\x18\x01 \x01(\x02\x1a\x42\n\x11\x42\x65rnoulliEstimate\x12-\n\x1e\x61\x64\x64itive_smoothing_pseudocount\x18\x04 \x01(\x02:\x05\x30.001\"\xbf\x05\n\x0fSimAgentMetrics\x12\x13\n\x0bscenario_id\x18\x01 \x01(\t\x12\x12\n\nmetametric\x18\x02 \x01(\x02\x12\"\n\x1a\x61verage_displacement_error\x18\x03 \x01(\x02\x12&\n\x1emin_average_displacement_error\x18\x13 \x01(\x02\x12\x1f\n\x17linear_speed_likelihood\x18\x04 \x01(\x02\x12&\n\x1elinear_acceleration_likelihood\x18\x05 \x01(\x02\x12 \n\x18\x61ngular_speed_likelihood\x18\x06 \x01(\x02\x12\'\n\x1f\x61ngular_acceleration_likelihood\x18\x07 \x01(\x02\x12-\n%distance_to_nearest_object_likelihood\x18\x08 \x01(\x02\x12\'\n\x1f\x63ollision_indication_likelihood\x18\t \x01(\x02\x12$\n\x1ctime_to_collision_likelihood\x18\n \x01(\x02\x12(\n distance_to_road_edge_likelihood\x18\x0b \x01(\x02\x12%\n\x1doffroad_indication_likelihood\x18\x0c \x01(\x02\x12 \n\x18num_placement_likelihood\x18\r \x01(\x02\x12!\n\x19num_removement_likelihood\x18\x0e \x01(\x02\x12%\n\x1d\x64istance_placement_likelihood\x18\x0f \x01(\x02\x12&\n\x1e\x64istance_removement_likelihood\x18\x10 \x01(\x02\x12 \n\x18simulated_collision_rate\x18\x11 \x01(\x02\x12\x1e\n\x16simulated_offroad_rate\x18\x12 \x01(\x02\"\xfe\x01\n\x18SimAgentsBucketedMetrics\x12\x1b\n\x13realism_meta_metric\x18\x01 \x01(\x02\x12\x19\n\x11kinematic_metrics\x18\x02 \x01(\x02\x12\x1b\n\x13interactive_metrics\x18\x05 \x01(\x02\x12\x19\n\x11map_based_metrics\x18\x06 \x01(\x02\x12\x1f\n\x17placement_based_metrics\x18\x07 \x01(\x02\x12\x0f\n\x07min_ade\x18\x08 \x01(\x02\x12 \n\x18simulated_collision_rate\x18\t \x01(\x02\x12\x1e\n\x16simulated_offroad_rate\x18\n \x01(\x02'
+)
+
+
+
+
+_SIMAGENTMETRICSCONFIG_FEATURECONFIG = _descriptor.Descriptor(
+  name='FeatureConfig',
+  full_name='long_metric.SimAgentMetricsConfig.FeatureConfig',
+  filename=None,
+  file=DESCRIPTOR,
+  containing_type=None,
+  create_key=_descriptor._internal_create_key,
+  fields=[
+    _descriptor.FieldDescriptor(
+      name='histogram', full_name='long_metric.SimAgentMetricsConfig.FeatureConfig.histogram', index=0,
+      number=1, type=11, cpp_type=10, label=1,
+      has_default_value=False, default_value=None,
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR,  create_key=_descriptor._internal_create_key),
+    _descriptor.FieldDescriptor(
+      name='kernel_density', full_name='long_metric.SimAgentMetricsConfig.FeatureConfig.kernel_density', index=1,
+      number=2, type=11, cpp_type=10, label=1,
+      has_default_value=False, default_value=None,
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR,  create_key=_descriptor._internal_create_key),
+    _descriptor.FieldDescriptor(
+      name='bernoulli', full_name='long_metric.SimAgentMetricsConfig.FeatureConfig.bernoulli', index=2,
+      number=3, type=11, cpp_type=10, label=1,
+      has_default_value=False, default_value=None,
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR,  create_key=_descriptor._internal_create_key),
+    _descriptor.FieldDescriptor(
+      name='independent_timesteps', full_name='long_metric.SimAgentMetricsConfig.FeatureConfig.independent_timesteps', index=3,
+      number=4, type=8, cpp_type=7, label=1,
+      has_default_value=False, default_value=False,
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR,  create_key=_descriptor._internal_create_key),
+    _descriptor.FieldDescriptor(
+      name='metametric_weight', full_name='long_metric.SimAgentMetricsConfig.FeatureConfig.metametric_weight', index=4,
+      number=5, type=2, cpp_type=6, label=1,
+      has_default_value=False, default_value=float(0),
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR,  create_key=_descriptor._internal_create_key),
+  ],
+  extensions=[
+  ],
+  nested_types=[],
+  enum_types=[
+  ],
+  serialized_options=None,
+  is_extendable=False,
+  syntax='proto2',
+  extension_ranges=[],
+  oneofs=[
+    _descriptor.OneofDescriptor(
+      name='estimator', full_name='long_metric.SimAgentMetricsConfig.FeatureConfig.estimator',
+      index=0, containing_type=None,
+      create_key=_descriptor._internal_create_key,
+    fields=[]),
+  ],
+  serialized_start=1091,
+  serialized_end=1411,
+)
+
+_SIMAGENTMETRICSCONFIG_HISTOGRAMESTIMATE = _descriptor.Descriptor(
+  name='HistogramEstimate',
+  full_name='long_metric.SimAgentMetricsConfig.HistogramEstimate',
+  filename=None,
+  file=DESCRIPTOR,
+  containing_type=None,
+  create_key=_descriptor._internal_create_key,
+  fields=[
+    _descriptor.FieldDescriptor(
+      name='min_val', full_name='long_metric.SimAgentMetricsConfig.HistogramEstimate.min_val', index=0,
+      number=1, type=2, cpp_type=6, label=1,
+      has_default_value=False, default_value=float(0),
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR,  create_key=_descriptor._internal_create_key),
+    _descriptor.FieldDescriptor(
+      name='max_val', full_name='long_metric.SimAgentMetricsConfig.HistogramEstimate.max_val', index=1,
+      number=2, type=2, cpp_type=6, label=1,
+      has_default_value=False, default_value=float(0),
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR,  create_key=_descriptor._internal_create_key),
+    _descriptor.FieldDescriptor(
+      name='num_bins', full_name='long_metric.SimAgentMetricsConfig.HistogramEstimate.num_bins', index=2,
+      number=3, type=5, cpp_type=1, label=1,
+      has_default_value=False, default_value=0,
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR,  create_key=_descriptor._internal_create_key),
+    _descriptor.FieldDescriptor(
+      name='additive_smoothing_pseudocount', full_name='long_metric.SimAgentMetricsConfig.HistogramEstimate.additive_smoothing_pseudocount', index=3,
+      number=4, type=2, cpp_type=6, label=1,
+      has_default_value=True, default_value=float(0.001),
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR,  create_key=_descriptor._internal_create_key),
+  ],
+  extensions=[
+  ],
+  nested_types=[],
+  enum_types=[
+  ],
+  serialized_options=None,
+  is_extendable=False,
+  syntax='proto2',
+  extension_ranges=[],
+  oneofs=[
+  ],
+  serialized_start=1413,
+  serialized_end=1531,
+)
+
+_SIMAGENTMETRICSCONFIG_KERNELDENSITYESTIMATE = _descriptor.Descriptor(
+  name='KernelDensityEstimate',
+  full_name='long_metric.SimAgentMetricsConfig.KernelDensityEstimate',
+  filename=None,
+  file=DESCRIPTOR,
+  containing_type=None,
+  create_key=_descriptor._internal_create_key,
+  fields=[
+    _descriptor.FieldDescriptor(
+      name='bandwidth', full_name='long_metric.SimAgentMetricsConfig.KernelDensityEstimate.bandwidth', index=0,
+      number=1, type=2, cpp_type=6, label=1,
+      has_default_value=False, default_value=float(0),
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR,  create_key=_descriptor._internal_create_key),
+  ],
+  extensions=[
+  ],
+  nested_types=[],
+  enum_types=[
+  ],
+  serialized_options=None,
+  is_extendable=False,
+  syntax='proto2',
+  extension_ranges=[],
+  oneofs=[
+  ],
+  serialized_start=1533,
+  serialized_end=1575,
+)
+
+_SIMAGENTMETRICSCONFIG_BERNOULLIESTIMATE = _descriptor.Descriptor(
+  name='BernoulliEstimate',
+  full_name='long_metric.SimAgentMetricsConfig.BernoulliEstimate',
+  filename=None,
+  file=DESCRIPTOR,
+  containing_type=None,
+  create_key=_descriptor._internal_create_key,
+  fields=[
+    _descriptor.FieldDescriptor(
+      name='additive_smoothing_pseudocount', full_name='long_metric.SimAgentMetricsConfig.BernoulliEstimate.additive_smoothing_pseudocount', index=0,
+      number=4, type=2, cpp_type=6, label=1,
+      has_default_value=True, default_value=float(0.001),
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR,  create_key=_descriptor._internal_create_key),
+  ],
+  extensions=[
+  ],
+  nested_types=[],
+  enum_types=[
+  ],
+  serialized_options=None,
+  is_extendable=False,
+  syntax='proto2',
+  extension_ranges=[],
+  oneofs=[
+  ],
+  serialized_start=1577,
+  serialized_end=1643,
+)
+
+_SIMAGENTMETRICSCONFIG = _descriptor.Descriptor(
+  name='SimAgentMetricsConfig',
+  full_name='long_metric.SimAgentMetricsConfig',
+  filename=None,
+  file=DESCRIPTOR,
+  containing_type=None,
+  create_key=_descriptor._internal_create_key,
+  fields=[
+    _descriptor.FieldDescriptor(
+      name='linear_speed', full_name='long_metric.SimAgentMetricsConfig.linear_speed', index=0,
+      number=1, type=11, cpp_type=10, label=1,
+      has_default_value=False, default_value=None,
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR,  create_key=_descriptor._internal_create_key),
+    _descriptor.FieldDescriptor(
+      name='linear_acceleration', full_name='long_metric.SimAgentMetricsConfig.linear_acceleration', index=1,
+      number=2, type=11, cpp_type=10, label=1,
+      has_default_value=False, default_value=None,
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR,  create_key=_descriptor._internal_create_key),
+    _descriptor.FieldDescriptor(
+      name='angular_speed', full_name='long_metric.SimAgentMetricsConfig.angular_speed', index=2,
+      number=3, type=11, cpp_type=10, label=1,
+      has_default_value=False, default_value=None,
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR,  create_key=_descriptor._internal_create_key),
+    _descriptor.FieldDescriptor(
+      name='angular_acceleration', full_name='long_metric.SimAgentMetricsConfig.angular_acceleration', index=3,
+      number=4, type=11, cpp_type=10, label=1,
+      has_default_value=False, default_value=None,
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR,  create_key=_descriptor._internal_create_key),
+    _descriptor.FieldDescriptor(
+      name='distance_to_nearest_object', full_name='long_metric.SimAgentMetricsConfig.distance_to_nearest_object', index=4,
+      number=5, type=11, cpp_type=10, label=1,
+      has_default_value=False, default_value=None,
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR,  create_key=_descriptor._internal_create_key),
+    _descriptor.FieldDescriptor(
+      name='collision_indication', full_name='long_metric.SimAgentMetricsConfig.collision_indication', index=5,
+      number=6, type=11, cpp_type=10, label=1,
+      has_default_value=False, default_value=None,
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR,  create_key=_descriptor._internal_create_key),
+    _descriptor.FieldDescriptor(
+      name='time_to_collision', full_name='long_metric.SimAgentMetricsConfig.time_to_collision', index=6,
+      number=7, type=11, cpp_type=10, label=1,
+      has_default_value=False, default_value=None,
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR,  create_key=_descriptor._internal_create_key),
+    _descriptor.FieldDescriptor(
+      name='distance_to_road_edge', full_name='long_metric.SimAgentMetricsConfig.distance_to_road_edge', index=7,
+      number=8, type=11, cpp_type=10, label=1,
+      has_default_value=False, default_value=None,
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR,  create_key=_descriptor._internal_create_key),
+    _descriptor.FieldDescriptor(
+      name='offroad_indication', full_name='long_metric.SimAgentMetricsConfig.offroad_indication', index=8,
+      number=9, type=11, cpp_type=10, label=1,
+      has_default_value=False, default_value=None,
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR,  create_key=_descriptor._internal_create_key),
+    _descriptor.FieldDescriptor(
+      name='num_placement', full_name='long_metric.SimAgentMetricsConfig.num_placement', index=9,
+      number=10, type=11, cpp_type=10, label=1,
+      has_default_value=False, default_value=None,
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR,  create_key=_descriptor._internal_create_key),
+    _descriptor.FieldDescriptor(
+      name='num_removement', full_name='long_metric.SimAgentMetricsConfig.num_removement', index=10,
+      number=11, type=11, cpp_type=10, label=1,
+      has_default_value=False, default_value=None,
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR,  create_key=_descriptor._internal_create_key),
+    _descriptor.FieldDescriptor(
+      name='distance_placement', full_name='long_metric.SimAgentMetricsConfig.distance_placement', index=11,
+      number=12, type=11, cpp_type=10, label=1,
+      has_default_value=False, default_value=None,
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR,  create_key=_descriptor._internal_create_key),
+    _descriptor.FieldDescriptor(
+      name='distance_removement', full_name='long_metric.SimAgentMetricsConfig.distance_removement', index=12,
+      number=13, type=11, cpp_type=10, label=1,
+      has_default_value=False, default_value=None,
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR,  create_key=_descriptor._internal_create_key),
+  ],
+  extensions=[
+  ],
+  nested_types=[_SIMAGENTMETRICSCONFIG_FEATURECONFIG, _SIMAGENTMETRICSCONFIG_HISTOGRAMESTIMATE, _SIMAGENTMETRICSCONFIG_KERNELDENSITYESTIMATE, _SIMAGENTMETRICSCONFIG_BERNOULLIESTIMATE, ],
+  enum_types=[
+  ],
+  serialized_options=None,
+  is_extendable=False,
+  syntax='proto2',
+  extension_ranges=[],
+  oneofs=[
+  ],
+  serialized_start=55,
+  serialized_end=1643,
+)
+
+
+_SIMAGENTMETRICS = _descriptor.Descriptor(
+  name='SimAgentMetrics',
+  full_name='long_metric.SimAgentMetrics',
+  filename=None,
+  file=DESCRIPTOR,
+  containing_type=None,
+  create_key=_descriptor._internal_create_key,
+  fields=[
+    _descriptor.FieldDescriptor(
+      name='scenario_id', full_name='long_metric.SimAgentMetrics.scenario_id', index=0,
+      number=1, type=9, cpp_type=9, label=1,
+      has_default_value=False, default_value=b"".decode('utf-8'),
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR,  create_key=_descriptor._internal_create_key),
+    _descriptor.FieldDescriptor(
+      name='metametric', full_name='long_metric.SimAgentMetrics.metametric', index=1,
+      number=2, type=2, cpp_type=6, label=1,
+      has_default_value=False, default_value=float(0),
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR,  create_key=_descriptor._internal_create_key),
+    _descriptor.FieldDescriptor(
+      name='average_displacement_error', full_name='long_metric.SimAgentMetrics.average_displacement_error', index=2,
+      number=3, type=2, cpp_type=6, label=1,
+      has_default_value=False, default_value=float(0),
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR,  create_key=_descriptor._internal_create_key),
+    _descriptor.FieldDescriptor(
+      name='min_average_displacement_error', full_name='long_metric.SimAgentMetrics.min_average_displacement_error', index=3,
+      number=19, type=2, cpp_type=6, label=1,
+      has_default_value=False, default_value=float(0),
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR,  create_key=_descriptor._internal_create_key),
+    _descriptor.FieldDescriptor(
+      name='linear_speed_likelihood', full_name='long_metric.SimAgentMetrics.linear_speed_likelihood', index=4,
+      number=4, type=2, cpp_type=6, label=1,
+      has_default_value=False, default_value=float(0),
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR,  create_key=_descriptor._internal_create_key),
+    _descriptor.FieldDescriptor(
+      name='linear_acceleration_likelihood', full_name='long_metric.SimAgentMetrics.linear_acceleration_likelihood', index=5,
+      number=5, type=2, cpp_type=6, label=1,
+      has_default_value=False, default_value=float(0),
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR,  create_key=_descriptor._internal_create_key),
+    _descriptor.FieldDescriptor(
+      name='angular_speed_likelihood', full_name='long_metric.SimAgentMetrics.angular_speed_likelihood', index=6,
+      number=6, type=2, cpp_type=6, label=1,
+      has_default_value=False, default_value=float(0),
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR,  create_key=_descriptor._internal_create_key),
+    _descriptor.FieldDescriptor(
+      name='angular_acceleration_likelihood', full_name='long_metric.SimAgentMetrics.angular_acceleration_likelihood', index=7,
+      number=7, type=2, cpp_type=6, label=1,
+      has_default_value=False, default_value=float(0),
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR,  create_key=_descriptor._internal_create_key),
+    _descriptor.FieldDescriptor(
+      name='distance_to_nearest_object_likelihood', full_name='long_metric.SimAgentMetrics.distance_to_nearest_object_likelihood', index=8,
+      number=8, type=2, cpp_type=6, label=1,
+      has_default_value=False, default_value=float(0),
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR,  create_key=_descriptor._internal_create_key),
+    _descriptor.FieldDescriptor(
+      name='collision_indication_likelihood', full_name='long_metric.SimAgentMetrics.collision_indication_likelihood', index=9,
+      number=9, type=2, cpp_type=6, label=1,
+      has_default_value=False, default_value=float(0),
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR,  create_key=_descriptor._internal_create_key),
+    _descriptor.FieldDescriptor(
+      name='time_to_collision_likelihood', full_name='long_metric.SimAgentMetrics.time_to_collision_likelihood', index=10,
+      number=10, type=2, cpp_type=6, label=1,
+      has_default_value=False, default_value=float(0),
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR,  create_key=_descriptor._internal_create_key),
+    _descriptor.FieldDescriptor(
+      name='distance_to_road_edge_likelihood', full_name='long_metric.SimAgentMetrics.distance_to_road_edge_likelihood', index=11,
+      number=11, type=2, cpp_type=6, label=1,
+      has_default_value=False, default_value=float(0),
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR,  create_key=_descriptor._internal_create_key),
+    _descriptor.FieldDescriptor(
+      name='offroad_indication_likelihood', full_name='long_metric.SimAgentMetrics.offroad_indication_likelihood', index=12,
+      number=12, type=2, cpp_type=6, label=1,
+      has_default_value=False, default_value=float(0),
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR,  create_key=_descriptor._internal_create_key),
+    _descriptor.FieldDescriptor(
+      name='num_placement_likelihood', full_name='long_metric.SimAgentMetrics.num_placement_likelihood', index=13,
+      number=13, type=2, cpp_type=6, label=1,
+      has_default_value=False, default_value=float(0),
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR,  create_key=_descriptor._internal_create_key),
+    _descriptor.FieldDescriptor(
+      name='num_removement_likelihood', full_name='long_metric.SimAgentMetrics.num_removement_likelihood', index=14,
+      number=14, type=2, cpp_type=6, label=1,
+      has_default_value=False, default_value=float(0),
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR,  create_key=_descriptor._internal_create_key),
+    _descriptor.FieldDescriptor(
+      name='distance_placement_likelihood', full_name='long_metric.SimAgentMetrics.distance_placement_likelihood', index=15,
+      number=15, type=2, cpp_type=6, label=1,
+      has_default_value=False, default_value=float(0),
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR,  create_key=_descriptor._internal_create_key),
+    _descriptor.FieldDescriptor(
+      name='distance_removement_likelihood', full_name='long_metric.SimAgentMetrics.distance_removement_likelihood', index=16,
+      number=16, type=2, cpp_type=6, label=1,
+      has_default_value=False, default_value=float(0),
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR,  create_key=_descriptor._internal_create_key),
+    _descriptor.FieldDescriptor(
+      name='simulated_collision_rate', full_name='long_metric.SimAgentMetrics.simulated_collision_rate', index=17,
+      number=17, type=2, cpp_type=6, label=1,
+      has_default_value=False, default_value=float(0),
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR,  create_key=_descriptor._internal_create_key),
+    _descriptor.FieldDescriptor(
+      name='simulated_offroad_rate', full_name='long_metric.SimAgentMetrics.simulated_offroad_rate', index=18,
+      number=18, type=2, cpp_type=6, label=1,
+      has_default_value=False, default_value=float(0),
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR,  create_key=_descriptor._internal_create_key),
+  ],
+  extensions=[
+  ],
+  nested_types=[],
+  enum_types=[
+  ],
+  serialized_options=None,
+  is_extendable=False,
+  syntax='proto2',
+  extension_ranges=[],
+  oneofs=[
+  ],
+  serialized_start=1646,
+  serialized_end=2349,
+)
+
+
+_SIMAGENTSBUCKETEDMETRICS = _descriptor.Descriptor(
+  name='SimAgentsBucketedMetrics',
+  full_name='long_metric.SimAgentsBucketedMetrics',
+  filename=None,
+  file=DESCRIPTOR,
+  containing_type=None,
+  create_key=_descriptor._internal_create_key,
+  fields=[
+    _descriptor.FieldDescriptor(
+      name='realism_meta_metric', full_name='long_metric.SimAgentsBucketedMetrics.realism_meta_metric', index=0,
+      number=1, type=2, cpp_type=6, label=1,
+      has_default_value=False, default_value=float(0),
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR,  create_key=_descriptor._internal_create_key),
+    _descriptor.FieldDescriptor(
+      name='kinematic_metrics', full_name='long_metric.SimAgentsBucketedMetrics.kinematic_metrics', index=1,
+      number=2, type=2, cpp_type=6, label=1,
+      has_default_value=False, default_value=float(0),
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR,  create_key=_descriptor._internal_create_key),
+    _descriptor.FieldDescriptor(
+      name='interactive_metrics', full_name='long_metric.SimAgentsBucketedMetrics.interactive_metrics', index=2,
+      number=5, type=2, cpp_type=6, label=1,
+      has_default_value=False, default_value=float(0),
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR,  create_key=_descriptor._internal_create_key),
+    _descriptor.FieldDescriptor(
+      name='map_based_metrics', full_name='long_metric.SimAgentsBucketedMetrics.map_based_metrics', index=3,
+      number=6, type=2, cpp_type=6, label=1,
+      has_default_value=False, default_value=float(0),
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR,  create_key=_descriptor._internal_create_key),
+    _descriptor.FieldDescriptor(
+      name='placement_based_metrics', full_name='long_metric.SimAgentsBucketedMetrics.placement_based_metrics', index=4,
+      number=7, type=2, cpp_type=6, label=1,
+      has_default_value=False, default_value=float(0),
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR,  create_key=_descriptor._internal_create_key),
+    _descriptor.FieldDescriptor(
+      name='min_ade', full_name='long_metric.SimAgentsBucketedMetrics.min_ade', index=5,
+      number=8, type=2, cpp_type=6, label=1,
+      has_default_value=False, default_value=float(0),
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR,  create_key=_descriptor._internal_create_key),
+    _descriptor.FieldDescriptor(
+      name='simulated_collision_rate', full_name='long_metric.SimAgentsBucketedMetrics.simulated_collision_rate', index=6,
+      number=9, type=2, cpp_type=6, label=1,
+      has_default_value=False, default_value=float(0),
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR,  create_key=_descriptor._internal_create_key),
+    _descriptor.FieldDescriptor(
+      name='simulated_offroad_rate', full_name='long_metric.SimAgentsBucketedMetrics.simulated_offroad_rate', index=7,
+      number=10, type=2, cpp_type=6, label=1,
+      has_default_value=False, default_value=float(0),
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR,  create_key=_descriptor._internal_create_key),
+  ],
+  extensions=[
+  ],
+  nested_types=[],
+  enum_types=[
+  ],
+  serialized_options=None,
+  is_extendable=False,
+  syntax='proto2',
+  extension_ranges=[],
+  oneofs=[
+  ],
+  serialized_start=2352,
+  serialized_end=2606,
+)
+
+_SIMAGENTMETRICSCONFIG_FEATURECONFIG.fields_by_name['histogram'].message_type = _SIMAGENTMETRICSCONFIG_HISTOGRAMESTIMATE
+_SIMAGENTMETRICSCONFIG_FEATURECONFIG.fields_by_name['kernel_density'].message_type = _SIMAGENTMETRICSCONFIG_KERNELDENSITYESTIMATE
+_SIMAGENTMETRICSCONFIG_FEATURECONFIG.fields_by_name['bernoulli'].message_type = _SIMAGENTMETRICSCONFIG_BERNOULLIESTIMATE
+_SIMAGENTMETRICSCONFIG_FEATURECONFIG.containing_type = _SIMAGENTMETRICSCONFIG
+_SIMAGENTMETRICSCONFIG_FEATURECONFIG.oneofs_by_name['estimator'].fields.append(
+  _SIMAGENTMETRICSCONFIG_FEATURECONFIG.fields_by_name['histogram'])
+_SIMAGENTMETRICSCONFIG_FEATURECONFIG.fields_by_name['histogram'].containing_oneof = _SIMAGENTMETRICSCONFIG_FEATURECONFIG.oneofs_by_name['estimator']
+_SIMAGENTMETRICSCONFIG_FEATURECONFIG.oneofs_by_name['estimator'].fields.append(
+  _SIMAGENTMETRICSCONFIG_FEATURECONFIG.fields_by_name['kernel_density'])
+_SIMAGENTMETRICSCONFIG_FEATURECONFIG.fields_by_name['kernel_density'].containing_oneof = _SIMAGENTMETRICSCONFIG_FEATURECONFIG.oneofs_by_name['estimator']
+_SIMAGENTMETRICSCONFIG_FEATURECONFIG.oneofs_by_name['estimator'].fields.append(
+  _SIMAGENTMETRICSCONFIG_FEATURECONFIG.fields_by_name['bernoulli'])
+_SIMAGENTMETRICSCONFIG_FEATURECONFIG.fields_by_name['bernoulli'].containing_oneof = _SIMAGENTMETRICSCONFIG_FEATURECONFIG.oneofs_by_name['estimator']
+_SIMAGENTMETRICSCONFIG_HISTOGRAMESTIMATE.containing_type = _SIMAGENTMETRICSCONFIG
+_SIMAGENTMETRICSCONFIG_KERNELDENSITYESTIMATE.containing_type = _SIMAGENTMETRICSCONFIG
+_SIMAGENTMETRICSCONFIG_BERNOULLIESTIMATE.containing_type = _SIMAGENTMETRICSCONFIG
+_SIMAGENTMETRICSCONFIG.fields_by_name['linear_speed'].message_type = _SIMAGENTMETRICSCONFIG_FEATURECONFIG
+_SIMAGENTMETRICSCONFIG.fields_by_name['linear_acceleration'].message_type = _SIMAGENTMETRICSCONFIG_FEATURECONFIG
+_SIMAGENTMETRICSCONFIG.fields_by_name['angular_speed'].message_type = _SIMAGENTMETRICSCONFIG_FEATURECONFIG
+_SIMAGENTMETRICSCONFIG.fields_by_name['angular_acceleration'].message_type = _SIMAGENTMETRICSCONFIG_FEATURECONFIG
+_SIMAGENTMETRICSCONFIG.fields_by_name['distance_to_nearest_object'].message_type = _SIMAGENTMETRICSCONFIG_FEATURECONFIG
+_SIMAGENTMETRICSCONFIG.fields_by_name['collision_indication'].message_type = _SIMAGENTMETRICSCONFIG_FEATURECONFIG
+_SIMAGENTMETRICSCONFIG.fields_by_name['time_to_collision'].message_type = _SIMAGENTMETRICSCONFIG_FEATURECONFIG
+_SIMAGENTMETRICSCONFIG.fields_by_name['distance_to_road_edge'].message_type = _SIMAGENTMETRICSCONFIG_FEATURECONFIG
+_SIMAGENTMETRICSCONFIG.fields_by_name['offroad_indication'].message_type = _SIMAGENTMETRICSCONFIG_FEATURECONFIG
+_SIMAGENTMETRICSCONFIG.fields_by_name['num_placement'].message_type = _SIMAGENTMETRICSCONFIG_FEATURECONFIG
+_SIMAGENTMETRICSCONFIG.fields_by_name['num_removement'].message_type = _SIMAGENTMETRICSCONFIG_FEATURECONFIG
+_SIMAGENTMETRICSCONFIG.fields_by_name['distance_placement'].message_type = _SIMAGENTMETRICSCONFIG_FEATURECONFIG
+_SIMAGENTMETRICSCONFIG.fields_by_name['distance_removement'].message_type = _SIMAGENTMETRICSCONFIG_FEATURECONFIG
+DESCRIPTOR.message_types_by_name['SimAgentMetricsConfig'] = _SIMAGENTMETRICSCONFIG
+DESCRIPTOR.message_types_by_name['SimAgentMetrics'] = _SIMAGENTMETRICS
+DESCRIPTOR.message_types_by_name['SimAgentsBucketedMetrics'] = _SIMAGENTSBUCKETEDMETRICS
+_sym_db.RegisterFileDescriptor(DESCRIPTOR)
+
+SimAgentMetricsConfig = _reflection.GeneratedProtocolMessageType('SimAgentMetricsConfig', (_message.Message,), {
+
+  'FeatureConfig' : _reflection.GeneratedProtocolMessageType('FeatureConfig', (_message.Message,), {
+    'DESCRIPTOR' : _SIMAGENTMETRICSCONFIG_FEATURECONFIG,
+    '__module__' : 'dev.metrics.protos.long_metrics_pb2'
+    # @@protoc_insertion_point(class_scope:long_metric.SimAgentMetricsConfig.FeatureConfig)
+    })
+  ,
+
+  'HistogramEstimate' : _reflection.GeneratedProtocolMessageType('HistogramEstimate', (_message.Message,), {
+    'DESCRIPTOR' : _SIMAGENTMETRICSCONFIG_HISTOGRAMESTIMATE,
+    '__module__' : 'dev.metrics.protos.long_metrics_pb2'
+    # @@protoc_insertion_point(class_scope:long_metric.SimAgentMetricsConfig.HistogramEstimate)
+    })
+  ,
+
+  'KernelDensityEstimate' : _reflection.GeneratedProtocolMessageType('KernelDensityEstimate', (_message.Message,), {
+    'DESCRIPTOR' : _SIMAGENTMETRICSCONFIG_KERNELDENSITYESTIMATE,
+    '__module__' : 'dev.metrics.protos.long_metrics_pb2'
+    # @@protoc_insertion_point(class_scope:long_metric.SimAgentMetricsConfig.KernelDensityEstimate)
+    })
+  ,
+
+  'BernoulliEstimate' : _reflection.GeneratedProtocolMessageType('BernoulliEstimate', (_message.Message,), {
+    'DESCRIPTOR' : _SIMAGENTMETRICSCONFIG_BERNOULLIESTIMATE,
+    '__module__' : 'dev.metrics.protos.long_metrics_pb2'
+    # @@protoc_insertion_point(class_scope:long_metric.SimAgentMetricsConfig.BernoulliEstimate)
+    })
+  ,
+  'DESCRIPTOR' : _SIMAGENTMETRICSCONFIG,
+  '__module__' : 'dev.metrics.protos.long_metrics_pb2'
+  # @@protoc_insertion_point(class_scope:long_metric.SimAgentMetricsConfig)
+  })
+_sym_db.RegisterMessage(SimAgentMetricsConfig)
+_sym_db.RegisterMessage(SimAgentMetricsConfig.FeatureConfig)
+_sym_db.RegisterMessage(SimAgentMetricsConfig.HistogramEstimate)
+_sym_db.RegisterMessage(SimAgentMetricsConfig.KernelDensityEstimate)
+_sym_db.RegisterMessage(SimAgentMetricsConfig.BernoulliEstimate)
+
+SimAgentMetrics = _reflection.GeneratedProtocolMessageType('SimAgentMetrics', (_message.Message,), {
+  'DESCRIPTOR' : _SIMAGENTMETRICS,
+  '__module__' : 'dev.metrics.protos.long_metrics_pb2'
+  # @@protoc_insertion_point(class_scope:long_metric.SimAgentMetrics)
+  })
+_sym_db.RegisterMessage(SimAgentMetrics)
+
+SimAgentsBucketedMetrics = _reflection.GeneratedProtocolMessageType('SimAgentsBucketedMetrics', (_message.Message,), {
+  'DESCRIPTOR' : _SIMAGENTSBUCKETEDMETRICS,
+  '__module__' : 'dev.metrics.protos.long_metrics_pb2'
+  # @@protoc_insertion_point(class_scope:long_metric.SimAgentsBucketedMetrics)
+  })
+_sym_db.RegisterMessage(SimAgentsBucketedMetrics)
+
+
+# @@protoc_insertion_point(module_scope)

seq_10_150_3_3_encode_occ_separate_offsets_bs8_128_no_seqindex_long_lastvalid_ablation_grid/backups/dev/metrics/protos/map_pb2.py

@@ -0,0 +1,1070 @@
+# -*- coding: utf-8 -*-
+# Generated by the protocol buffer compiler.  DO NOT EDIT!
+# source: dev/metrics/protos/map.proto
+
+from google.protobuf import descriptor as _descriptor
+from google.protobuf import message as _message
+from google.protobuf import reflection as _reflection
+from google.protobuf import symbol_database as _symbol_database
+# @@protoc_insertion_point(imports)
+
+_sym_db = _symbol_database.Default()
+
+
+
+
+DESCRIPTOR = _descriptor.FileDescriptor(
+  name='dev/metrics/protos/map.proto',
+  package='long_metric',
+  syntax='proto2',
+  serialized_options=None,
+  create_key=_descriptor._internal_create_key,
+  serialized_pb=b'\n\x1c\x64\x65v/metrics/protos/map.proto\x12\x0blong_metric\"g\n\x03Map\x12-\n\x0cmap_features\x18\x01 \x03(\x0b\x32\x17.long_metric.MapFeature\x12\x31\n\x0e\x64ynamic_states\x18\x02 \x03(\x0b\x32\x19.long_metric.DynamicState\"c\n\x0c\x44ynamicState\x12\x19\n\x11timestamp_seconds\x18\x01 \x01(\x01\x12\x38\n\x0blane_states\x18\x02 \x03(\x0b\x32#.long_metric.TrafficSignalLaneState\"\xfe\x02\n\x16TrafficSignalLaneState\x12\x0c\n\x04lane\x18\x01 \x01(\x03\x12\x38\n\x05state\x18\x02 \x01(\x0e\x32).long_metric.TrafficSignalLaneState.State\x12)\n\nstop_point\x18\x03 \x01(\x0b\x32\x15.long_metric.MapPoint\"\xf0\x01\n\x05State\x12\x16\n\x12LANE_STATE_UNKNOWN\x10\x00\x12\x19\n\x15LANE_STATE_ARROW_STOP\x10\x01\x12\x1c\n\x18LANE_STATE_ARROW_CAUTION\x10\x02\x12\x17\n\x13LANE_STATE_ARROW_GO\x10\x03\x12\x13\n\x0fLANE_STATE_STOP\x10\x04\x12\x16\n\x12LANE_STATE_CAUTION\x10\x05\x12\x11\n\rLANE_STATE_GO\x10\x06\x12\x1c\n\x18LANE_STATE_FLASHING_STOP\x10\x07\x12\x1f\n\x1bLANE_STATE_FLASHING_CAUTION\x10\x08\"\xdb\x02\n\nMapFeature\x12\n\n\x02id\x18\x01 \x01(\x03\x12\'\n\x04lane\x18\x03 \x01(\x0b\x32\x17.long_metric.LaneCenterH\x00\x12*\n\troad_line\x18\x04 \x01(\x0b\x32\x15.long_metric.RoadLineH\x00\x12*\n\troad_edge\x18\x05 \x01(\x0b\x32\x15.long_metric.RoadEdgeH\x00\x12*\n\tstop_sign\x18\x07 \x01(\x0b\x32\x15.long_metric.StopSignH\x00\x12+\n\tcrosswalk\x18\x08 \x01(\x0b\x32\x16.long_metric.CrosswalkH\x00\x12,\n\nspeed_bump\x18\t \x01(\x0b\x32\x16.long_metric.SpeedBumpH\x00\x12)\n\x08\x64riveway\x18\n \x01(\x0b\x32\x15.long_metric.DrivewayH\x00\x42\x0e\n\x0c\x66\x65\x61ture_data\"+\n\x08MapPoint\x12\t\n\x01x\x18\x01 \x01(\x01\x12\t\n\x01y\x18\x02 \x01(\x01\x12\t\n\x01z\x18\x03 \x01(\x01\"\x9b\x01\n\x0f\x42oundarySegment\x12\x18\n\x10lane_start_index\x18\x01 \x01(\x05\x12\x16\n\x0elane_end_index\x18\x02 \x01(\x05\x12\x1b\n\x13\x62oundary_feature_id\x18\x03 \x01(\x03\x12\x39\n\rboundary_type\x18\x04 \x01(\x0e\x32\".long_metric.RoadLine.RoadLineType\"\xc0\x01\n\x0cLaneNeighbor\x12\x12\n\nfeature_id\x18\x01 \x01(\x03\x12\x18\n\x10self_start_index\x18\x02 \x01(\x05\x12\x16\n\x0eself_end_index\x18\x03 \x01(\x05\x12\x1c\n\x14neighbor_start_index\x18\x04 \x01(\x05\x12\x1a\n\x12neighbor_end_index\x18\x05 \x01(\x05\x12\x30\n\nboundaries\x18\x06 \x03(\x0b\x32\x1c.long_metric.BoundarySegment\"\xfb\x03\n\nLaneCenter\x12\x17\n\x0fspeed_limit_mph\x18\x01 \x01(\x01\x12.\n\x04type\x18\x02 \x01(\x0e\x32 .long_metric.LaneCenter.LaneType\x12\x15\n\rinterpolating\x18\x03 \x01(\x08\x12\'\n\x08polyline\x18\x08 \x03(\x0b\x32\x15.long_metric.MapPoint\x12\x17\n\x0b\x65ntry_lanes\x18\t \x03(\x03\x42\x02\x10\x01\x12\x16\n\nexit_lanes\x18\n \x03(\x03\x42\x02\x10\x01\x12\x35\n\x0fleft_boundaries\x18\r \x03(\x0b\x32\x1c.long_metric.BoundarySegment\x12\x36\n\x10right_boundaries\x18\x0e \x03(\x0b\x32\x1c.long_metric.BoundarySegment\x12\x31\n\x0eleft_neighbors\x18\x0b \x03(\x0b\x32\x19.long_metric.LaneNeighbor\x12\x32\n\x0fright_neighbors\x18\x0c \x03(\x0b\x32\x19.long_metric.LaneNeighbor\"]\n\x08LaneType\x12\x12\n\x0eTYPE_UNDEFINED\x10\x00\x12\x10\n\x0cTYPE_FREEWAY\x10\x01\x12\x17\n\x13TYPE_SURFACE_STREET\x10\x02\x12\x12\n\x0eTYPE_BIKE_LANE\x10\x03\"\xbf\x01\n\x08RoadEdge\x12\x30\n\x04type\x18\x01 \x01(\x0e\x32\".long_metric.RoadEdge.RoadEdgeType\x12\'\n\x08polyline\x18\x02 \x03(\x0b\x32\x15.long_metric.MapPoint\"X\n\x0cRoadEdgeType\x12\x10\n\x0cTYPE_UNKNOWN\x10\x00\x12\x1b\n\x17TYPE_ROAD_EDGE_BOUNDARY\x10\x01\x12\x19\n\x15TYPE_ROAD_EDGE_MEDIAN\x10\x02\"\xfa\x02\n\x08RoadLine\x12\x30\n\x04type\x18\x01 \x01(\x0e\x32\".long_metric.RoadLine.RoadLineType\x12\'\n\x08polyline\x18\x02 \x03(\x0b\x32\x15.long_metric.MapPoint\"\x92\x02\n\x0cRoadLineType\x12\x10\n\x0cTYPE_UNKNOWN\x10\x00\x12\x1c\n\x18TYPE_BROKEN_SINGLE_WHITE\x10\x01\x12\x1b\n\x17TYPE_SOLID_SINGLE_WHITE\x10\x02\x12\x1b\n\x17TYPE_SOLID_DOUBLE_WHITE\x10\x03\x12\x1d\n\x19TYPE_BROKEN_SINGLE_YELLOW\x10\x04\x12\x1d\n\x19TYPE_BROKEN_DOUBLE_YELLOW\x10\x05\x12\x1c\n\x18TYPE_SOLID_SINGLE_YELLOW\x10\x06\x12\x1c\n\x18TYPE_SOLID_DOUBLE_YELLOW\x10\x07\x12\x1e\n\x1aTYPE_PASSING_DOUBLE_YELLOW\x10\x08\"A\n\x08StopSign\x12\x0c\n\x04lane\x18\x01 \x03(\x03\x12\'\n\x08position\x18\x02 \x01(\x0b\x32\x15.long_metric.MapPoint\"3\n\tCrosswalk\x12&\n\x07polygon\x18\x01 \x03(\x0b\x32\x15.long_metric.MapPoint\"3\n\tSpeedBump\x12&\n\x07polygon\x18\x01 \x03(\x0b\x32\x15.long_metric.MapPoint\"2\n\x08\x44riveway\x12&\n\x07polygon\x18\x01 \x03(\x0b\x32\x15.long_metric.MapPoint'
+)
+
+
+
+_TRAFFICSIGNALLANESTATE_STATE = _descriptor.EnumDescriptor(
+  name='State',
+  full_name='long_metric.TrafficSignalLaneState.State',
+  filename=None,
+  file=DESCRIPTOR,
+  create_key=_descriptor._internal_create_key,
+  values=[
+    _descriptor.EnumValueDescriptor(
+      name='LANE_STATE_UNKNOWN', index=0, number=0,
+      serialized_options=None,
+      type=None,
+      create_key=_descriptor._internal_create_key),
+    _descriptor.EnumValueDescriptor(
+      name='LANE_STATE_ARROW_STOP', index=1, number=1,
+      serialized_options=None,
+      type=None,
+      create_key=_descriptor._internal_create_key),
+    _descriptor.EnumValueDescriptor(
+      name='LANE_STATE_ARROW_CAUTION', index=2, number=2,
+      serialized_options=None,
+      type=None,
+      create_key=_descriptor._internal_create_key),
+    _descriptor.EnumValueDescriptor(
+      name='LANE_STATE_ARROW_GO', index=3, number=3,
+      serialized_options=None,
+      type=None,
+      create_key=_descriptor._internal_create_key),
+    _descriptor.EnumValueDescriptor(
+      name='LANE_STATE_STOP', index=4, number=4,
+      serialized_options=None,
+      type=None,
+      create_key=_descriptor._internal_create_key),
+    _descriptor.EnumValueDescriptor(
+      name='LANE_STATE_CAUTION', index=5, number=5,
+      serialized_options=None,
+      type=None,
+      create_key=_descriptor._internal_create_key),
+    _descriptor.EnumValueDescriptor(
+      name='LANE_STATE_GO', index=6, number=6,
+      serialized_options=None,
+      type=None,
+      create_key=_descriptor._internal_create_key),
+    _descriptor.EnumValueDescriptor(
+      name='LANE_STATE_FLASHING_STOP', index=7, number=7,
+      serialized_options=None,
+      type=None,
+      create_key=_descriptor._internal_create_key),
+    _descriptor.EnumValueDescriptor(
+      name='LANE_STATE_FLASHING_CAUTION', index=8, number=8,
+      serialized_options=None,
+      type=None,
+      create_key=_descriptor._internal_create_key),
+  ],
+  containing_type=None,
+  serialized_options=None,
+  serialized_start=394,
+  serialized_end=634,
+)
+_sym_db.RegisterEnumDescriptor(_TRAFFICSIGNALLANESTATE_STATE)
+
+_LANECENTER_LANETYPE = _descriptor.EnumDescriptor(
+  name='LaneType',
+  full_name='long_metric.LaneCenter.LaneType',
+  filename=None,
+  file=DESCRIPTOR,
+  create_key=_descriptor._internal_create_key,
+  values=[
+    _descriptor.EnumValueDescriptor(
+      name='TYPE_UNDEFINED', index=0, number=0,
+      serialized_options=None,
+      type=None,
+      create_key=_descriptor._internal_create_key),
+    _descriptor.EnumValueDescriptor(
+      name='TYPE_FREEWAY', index=1, number=1,
+      serialized_options=None,
+      type=None,
+      create_key=_descriptor._internal_create_key),
+    _descriptor.EnumValueDescriptor(
+      name='TYPE_SURFACE_STREET', index=2, number=2,
+      serialized_options=None,
+      type=None,
+      create_key=_descriptor._internal_create_key),
+    _descriptor.EnumValueDescriptor(
+      name='TYPE_BIKE_LANE', index=3, number=3,
+      serialized_options=None,
+      type=None,
+      create_key=_descriptor._internal_create_key),
+  ],
+  containing_type=None,
+  serialized_options=None,
+  serialized_start=1799,
+  serialized_end=1892,
+)
+_sym_db.RegisterEnumDescriptor(_LANECENTER_LANETYPE)
+
+_ROADEDGE_ROADEDGETYPE = _descriptor.EnumDescriptor(
+  name='RoadEdgeType',
+  full_name='long_metric.RoadEdge.RoadEdgeType',
+  filename=None,
+  file=DESCRIPTOR,
+  create_key=_descriptor._internal_create_key,
+  values=[
+    _descriptor.EnumValueDescriptor(
+      name='TYPE_UNKNOWN', index=0, number=0,
+      serialized_options=None,
+      type=None,
+      create_key=_descriptor._internal_create_key),
+    _descriptor.EnumValueDescriptor(
+      name='TYPE_ROAD_EDGE_BOUNDARY', index=1, number=1,
+      serialized_options=None,
+      type=None,
+      create_key=_descriptor._internal_create_key),
+    _descriptor.EnumValueDescriptor(
+      name='TYPE_ROAD_EDGE_MEDIAN', index=2, number=2,
+      serialized_options=None,
+      type=None,
+      create_key=_descriptor._internal_create_key),
+  ],
+  containing_type=None,
+  serialized_options=None,
+  serialized_start=1998,
+  serialized_end=2086,
+)
+_sym_db.RegisterEnumDescriptor(_ROADEDGE_ROADEDGETYPE)
+
+_ROADLINE_ROADLINETYPE = _descriptor.EnumDescriptor(
+  name='RoadLineType',
+  full_name='long_metric.RoadLine.RoadLineType',
+  filename=None,
+  file=DESCRIPTOR,
+  create_key=_descriptor._internal_create_key,
+  values=[
+    _descriptor.EnumValueDescriptor(
+      name='TYPE_UNKNOWN', index=0, number=0,
+      serialized_options=None,
+      type=None,
+      create_key=_descriptor._internal_create_key),
+    _descriptor.EnumValueDescriptor(
+      name='TYPE_BROKEN_SINGLE_WHITE', index=1, number=1,
+      serialized_options=None,
+      type=None,
+      create_key=_descriptor._internal_create_key),
+    _descriptor.EnumValueDescriptor(
+      name='TYPE_SOLID_SINGLE_WHITE', index=2, number=2,
+      serialized_options=None,
+      type=None,
+      create_key=_descriptor._internal_create_key),
+    _descriptor.EnumValueDescriptor(
+      name='TYPE_SOLID_DOUBLE_WHITE', index=3, number=3,
+      serialized_options=None,
+      type=None,
+      create_key=_descriptor._internal_create_key),
+    _descriptor.EnumValueDescriptor(
+      name='TYPE_BROKEN_SINGLE_YELLOW', index=4, number=4,
+      serialized_options=None,
+      type=None,
+      create_key=_descriptor._internal_create_key),
+    _descriptor.EnumValueDescriptor(
+      name='TYPE_BROKEN_DOUBLE_YELLOW', index=5, number=5,
+      serialized_options=None,
+      type=None,
+      create_key=_descriptor._internal_create_key),
+    _descriptor.EnumValueDescriptor(
+      name='TYPE_SOLID_SINGLE_YELLOW', index=6, number=6,
+      serialized_options=None,
+      type=None,
+      create_key=_descriptor._internal_create_key),
+    _descriptor.EnumValueDescriptor(
+      name='TYPE_SOLID_DOUBLE_YELLOW', index=7, number=7,
+      serialized_options=None,
+      type=None,
+      create_key=_descriptor._internal_create_key),
+    _descriptor.EnumValueDescriptor(
+      name='TYPE_PASSING_DOUBLE_YELLOW', index=8, number=8,
+      serialized_options=None,
+      type=None,
+      create_key=_descriptor._internal_create_key),
+  ],
+  containing_type=None,
+  serialized_options=None,
+  serialized_start=2193,
+  serialized_end=2467,
+)
+_sym_db.RegisterEnumDescriptor(_ROADLINE_ROADLINETYPE)
+
+
+_MAP = _descriptor.Descriptor(
+  name='Map',
+  full_name='long_metric.Map',
+  filename=None,
+  file=DESCRIPTOR,
+  containing_type=None,
+  create_key=_descriptor._internal_create_key,
+  fields=[
+    _descriptor.FieldDescriptor(
+      name='map_features', full_name='long_metric.Map.map_features', index=0,
+      number=1, type=11, cpp_type=10, label=3,
+      has_default_value=False, default_value=[],
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR,  create_key=_descriptor._internal_create_key),
+    _descriptor.FieldDescriptor(
+      name='dynamic_states', full_name='long_metric.Map.dynamic_states', index=1,
+      number=2, type=11, cpp_type=10, label=3,
+      has_default_value=False, default_value=[],
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR,  create_key=_descriptor._internal_create_key),
+  ],
+  extensions=[
+  ],
+  nested_types=[],
+  enum_types=[
+  ],
+  serialized_options=None,
+  is_extendable=False,
+  syntax='proto2',
+  extension_ranges=[],
+  oneofs=[
+  ],
+  serialized_start=45,
+  serialized_end=148,
+)
+
+
+_DYNAMICSTATE = _descriptor.Descriptor(
+  name='DynamicState',
+  full_name='long_metric.DynamicState',
+  filename=None,
+  file=DESCRIPTOR,
+  containing_type=None,
+  create_key=_descriptor._internal_create_key,
+  fields=[
+    _descriptor.FieldDescriptor(
+      name='timestamp_seconds', full_name='long_metric.DynamicState.timestamp_seconds', index=0,
+      number=1, type=1, cpp_type=5, label=1,
+      has_default_value=False, default_value=float(0),
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR,  create_key=_descriptor._internal_create_key),
+    _descriptor.FieldDescriptor(
+      name='lane_states', full_name='long_metric.DynamicState.lane_states', index=1,
+      number=2, type=11, cpp_type=10, label=3,
+      has_default_value=False, default_value=[],
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR,  create_key=_descriptor._internal_create_key),
+  ],
+  extensions=[
+  ],
+  nested_types=[],
+  enum_types=[
+  ],
+  serialized_options=None,
+  is_extendable=False,
+  syntax='proto2',
+  extension_ranges=[],
+  oneofs=[
+  ],
+  serialized_start=150,
+  serialized_end=249,
+)
+
+
+_TRAFFICSIGNALLANESTATE = _descriptor.Descriptor(
+  name='TrafficSignalLaneState',
+  full_name='long_metric.TrafficSignalLaneState',
+  filename=None,
+  file=DESCRIPTOR,
+  containing_type=None,
+  create_key=_descriptor._internal_create_key,
+  fields=[
+    _descriptor.FieldDescriptor(
+      name='lane', full_name='long_metric.TrafficSignalLaneState.lane', index=0,
+      number=1, type=3, cpp_type=2, label=1,
+      has_default_value=False, default_value=0,
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR,  create_key=_descriptor._internal_create_key),
+    _descriptor.FieldDescriptor(
+      name='state', full_name='long_metric.TrafficSignalLaneState.state', index=1,
+      number=2, type=14, cpp_type=8, label=1,
+      has_default_value=False, default_value=0,
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR,  create_key=_descriptor._internal_create_key),
+    _descriptor.FieldDescriptor(
+      name='stop_point', full_name='long_metric.TrafficSignalLaneState.stop_point', index=2,
+      number=3, type=11, cpp_type=10, label=1,
+      has_default_value=False, default_value=None,
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR,  create_key=_descriptor._internal_create_key),
+  ],
+  extensions=[
+  ],
+  nested_types=[],
+  enum_types=[
+    _TRAFFICSIGNALLANESTATE_STATE,
+  ],
+  serialized_options=None,
+  is_extendable=False,
+  syntax='proto2',
+  extension_ranges=[],
+  oneofs=[
+  ],
+  serialized_start=252,
+  serialized_end=634,
+)
+
+
+_MAPFEATURE = _descriptor.Descriptor(
+  name='MapFeature',
+  full_name='long_metric.MapFeature',
+  filename=None,
+  file=DESCRIPTOR,
+  containing_type=None,
+  create_key=_descriptor._internal_create_key,
+  fields=[
+    _descriptor.FieldDescriptor(
+      name='id', full_name='long_metric.MapFeature.id', index=0,
+      number=1, type=3, cpp_type=2, label=1,
+      has_default_value=False, default_value=0,
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR,  create_key=_descriptor._internal_create_key),
+    _descriptor.FieldDescriptor(
+      name='lane', full_name='long_metric.MapFeature.lane', index=1,
+      number=3, type=11, cpp_type=10, label=1,
+      has_default_value=False, default_value=None,
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR,  create_key=_descriptor._internal_create_key),
+    _descriptor.FieldDescriptor(
+      name='road_line', full_name='long_metric.MapFeature.road_line', index=2,
+      number=4, type=11, cpp_type=10, label=1,
+      has_default_value=False, default_value=None,
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR,  create_key=_descriptor._internal_create_key),
+    _descriptor.FieldDescriptor(
+      name='road_edge', full_name='long_metric.MapFeature.road_edge', index=3,
+      number=5, type=11, cpp_type=10, label=1,
+      has_default_value=False, default_value=None,
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR,  create_key=_descriptor._internal_create_key),
+    _descriptor.FieldDescriptor(
+      name='stop_sign', full_name='long_metric.MapFeature.stop_sign', index=4,
+      number=7, type=11, cpp_type=10, label=1,
+      has_default_value=False, default_value=None,
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR,  create_key=_descriptor._internal_create_key),
+    _descriptor.FieldDescriptor(
+      name='crosswalk', full_name='long_metric.MapFeature.crosswalk', index=5,
+      number=8, type=11, cpp_type=10, label=1,
+      has_default_value=False, default_value=None,
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR,  create_key=_descriptor._internal_create_key),
+    _descriptor.FieldDescriptor(
+      name='speed_bump', full_name='long_metric.MapFeature.speed_bump', index=6,
+      number=9, type=11, cpp_type=10, label=1,
+      has_default_value=False, default_value=None,
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR,  create_key=_descriptor._internal_create_key),
+    _descriptor.FieldDescriptor(
+      name='driveway', full_name='long_metric.MapFeature.driveway', index=7,
+      number=10, type=11, cpp_type=10, label=1,
+      has_default_value=False, default_value=None,
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR,  create_key=_descriptor._internal_create_key),
+  ],
+  extensions=[
+  ],
+  nested_types=[],
+  enum_types=[
+  ],
+  serialized_options=None,
+  is_extendable=False,
+  syntax='proto2',
+  extension_ranges=[],
+  oneofs=[
+    _descriptor.OneofDescriptor(
+      name='feature_data', full_name='long_metric.MapFeature.feature_data',
+      index=0, containing_type=None,
+      create_key=_descriptor._internal_create_key,
+    fields=[]),
+  ],
+  serialized_start=637,
+  serialized_end=984,
+)
+
+
+_MAPPOINT = _descriptor.Descriptor(
+  name='MapPoint',
+  full_name='long_metric.MapPoint',
+  filename=None,
+  file=DESCRIPTOR,
+  containing_type=None,
+  create_key=_descriptor._internal_create_key,
+  fields=[
+    _descriptor.FieldDescriptor(
+      name='x', full_name='long_metric.MapPoint.x', index=0,
+      number=1, type=1, cpp_type=5, label=1,
+      has_default_value=False, default_value=float(0),
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR,  create_key=_descriptor._internal_create_key),
+    _descriptor.FieldDescriptor(
+      name='y', full_name='long_metric.MapPoint.y', index=1,
+      number=2, type=1, cpp_type=5, label=1,
+      has_default_value=False, default_value=float(0),
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR,  create_key=_descriptor._internal_create_key),
+    _descriptor.FieldDescriptor(
+      name='z', full_name='long_metric.MapPoint.z', index=2,
+      number=3, type=1, cpp_type=5, label=1,
+      has_default_value=False, default_value=float(0),
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR,  create_key=_descriptor._internal_create_key),
+  ],
+  extensions=[
+  ],
+  nested_types=[],
+  enum_types=[
+  ],
+  serialized_options=None,
+  is_extendable=False,
+  syntax='proto2',
+  extension_ranges=[],
+  oneofs=[
+  ],
+  serialized_start=986,
+  serialized_end=1029,
+)
+
+
+_BOUNDARYSEGMENT = _descriptor.Descriptor(
+  name='BoundarySegment',
+  full_name='long_metric.BoundarySegment',
+  filename=None,
+  file=DESCRIPTOR,
+  containing_type=None,
+  create_key=_descriptor._internal_create_key,
+  fields=[
+    _descriptor.FieldDescriptor(
+      name='lane_start_index', full_name='long_metric.BoundarySegment.lane_start_index', index=0,
+      number=1, type=5, cpp_type=1, label=1,
+      has_default_value=False, default_value=0,
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR,  create_key=_descriptor._internal_create_key),
+    _descriptor.FieldDescriptor(
+      name='lane_end_index', full_name='long_metric.BoundarySegment.lane_end_index', index=1,
+      number=2, type=5, cpp_type=1, label=1,
+      has_default_value=False, default_value=0,
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR,  create_key=_descriptor._internal_create_key),
+    _descriptor.FieldDescriptor(
+      name='boundary_feature_id', full_name='long_metric.BoundarySegment.boundary_feature_id', index=2,
+      number=3, type=3, cpp_type=2, label=1,
+      has_default_value=False, default_value=0,
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR,  create_key=_descriptor._internal_create_key),
+    _descriptor.FieldDescriptor(
+      name='boundary_type', full_name='long_metric.BoundarySegment.boundary_type', index=3,
+      number=4, type=14, cpp_type=8, label=1,
+      has_default_value=False, default_value=0,
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR,  create_key=_descriptor._internal_create_key),
+  ],
+  extensions=[
+  ],
+  nested_types=[],
+  enum_types=[
+  ],
+  serialized_options=None,
+  is_extendable=False,
+  syntax='proto2',
+  extension_ranges=[],
+  oneofs=[
+  ],
+  serialized_start=1032,
+  serialized_end=1187,
+)
+
+
+_LANENEIGHBOR = _descriptor.Descriptor(
+  name='LaneNeighbor',
+  full_name='long_metric.LaneNeighbor',
+  filename=None,
+  file=DESCRIPTOR,
+  containing_type=None,
+  create_key=_descriptor._internal_create_key,
+  fields=[
+    _descriptor.FieldDescriptor(
+      name='feature_id', full_name='long_metric.LaneNeighbor.feature_id', index=0,
+      number=1, type=3, cpp_type=2, label=1,
+      has_default_value=False, default_value=0,
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR,  create_key=_descriptor._internal_create_key),
+    _descriptor.FieldDescriptor(
+      name='self_start_index', full_name='long_metric.LaneNeighbor.self_start_index', index=1,
+      number=2, type=5, cpp_type=1, label=1,
+      has_default_value=False, default_value=0,
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR,  create_key=_descriptor._internal_create_key),
+    _descriptor.FieldDescriptor(
+      name='self_end_index', full_name='long_metric.LaneNeighbor.self_end_index', index=2,
+      number=3, type=5, cpp_type=1, label=1,
+      has_default_value=False, default_value=0,
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR,  create_key=_descriptor._internal_create_key),
+    _descriptor.FieldDescriptor(
+      name='neighbor_start_index', full_name='long_metric.LaneNeighbor.neighbor_start_index', index=3,
+      number=4, type=5, cpp_type=1, label=1,
+      has_default_value=False, default_value=0,
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR,  create_key=_descriptor._internal_create_key),
+    _descriptor.FieldDescriptor(
+      name='neighbor_end_index', full_name='long_metric.LaneNeighbor.neighbor_end_index', index=4,
+      number=5, type=5, cpp_type=1, label=1,
+      has_default_value=False, default_value=0,
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR,  create_key=_descriptor._internal_create_key),
+    _descriptor.FieldDescriptor(
+      name='boundaries', full_name='long_metric.LaneNeighbor.boundaries', index=5,
+      number=6, type=11, cpp_type=10, label=3,
+      has_default_value=False, default_value=[],
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR,  create_key=_descriptor._internal_create_key),
+  ],
+  extensions=[
+  ],
+  nested_types=[],
+  enum_types=[
+  ],
+  serialized_options=None,
+  is_extendable=False,
+  syntax='proto2',
+  extension_ranges=[],
+  oneofs=[
+  ],
+  serialized_start=1190,
+  serialized_end=1382,
+)
+
+
+_LANECENTER = _descriptor.Descriptor(
+  name='LaneCenter',
+  full_name='long_metric.LaneCenter',
+  filename=None,
+  file=DESCRIPTOR,
+  containing_type=None,
+  create_key=_descriptor._internal_create_key,
+  fields=[
+    _descriptor.FieldDescriptor(
+      name='speed_limit_mph', full_name='long_metric.LaneCenter.speed_limit_mph', index=0,
+      number=1, type=1, cpp_type=5, label=1,
+      has_default_value=False, default_value=float(0),
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR,  create_key=_descriptor._internal_create_key),
+    _descriptor.FieldDescriptor(
+      name='type', full_name='long_metric.LaneCenter.type', index=1,
+      number=2, type=14, cpp_type=8, label=1,
+      has_default_value=False, default_value=0,
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR,  create_key=_descriptor._internal_create_key),
+    _descriptor.FieldDescriptor(
+      name='interpolating', full_name='long_metric.LaneCenter.interpolating', index=2,
+      number=3, type=8, cpp_type=7, label=1,
+      has_default_value=False, default_value=False,
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR,  create_key=_descriptor._internal_create_key),
+    _descriptor.FieldDescriptor(
+      name='polyline', full_name='long_metric.LaneCenter.polyline', index=3,
+      number=8, type=11, cpp_type=10, label=3,
+      has_default_value=False, default_value=[],
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR,  create_key=_descriptor._internal_create_key),
+    _descriptor.FieldDescriptor(
+      name='entry_lanes', full_name='long_metric.LaneCenter.entry_lanes', index=4,
+      number=9, type=3, cpp_type=2, label=3,
+      has_default_value=False, default_value=[],
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=b'\020\001', file=DESCRIPTOR,  create_key=_descriptor._internal_create_key),
+    _descriptor.FieldDescriptor(
+      name='exit_lanes', full_name='long_metric.LaneCenter.exit_lanes', index=5,
+      number=10, type=3, cpp_type=2, label=3,
+      has_default_value=False, default_value=[],
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=b'\020\001', file=DESCRIPTOR,  create_key=_descriptor._internal_create_key),
+    _descriptor.FieldDescriptor(
+      name='left_boundaries', full_name='long_metric.LaneCenter.left_boundaries', index=6,
+      number=13, type=11, cpp_type=10, label=3,
+      has_default_value=False, default_value=[],
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR,  create_key=_descriptor._internal_create_key),
+    _descriptor.FieldDescriptor(
+      name='right_boundaries', full_name='long_metric.LaneCenter.right_boundaries', index=7,
+      number=14, type=11, cpp_type=10, label=3,
+      has_default_value=False, default_value=[],
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR,  create_key=_descriptor._internal_create_key),
+    _descriptor.FieldDescriptor(
+      name='left_neighbors', full_name='long_metric.LaneCenter.left_neighbors', index=8,
+      number=11, type=11, cpp_type=10, label=3,
+      has_default_value=False, default_value=[],
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR,  create_key=_descriptor._internal_create_key),
+    _descriptor.FieldDescriptor(
+      name='right_neighbors', full_name='long_metric.LaneCenter.right_neighbors', index=9,
+      number=12, type=11, cpp_type=10, label=3,
+      has_default_value=False, default_value=[],
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR,  create_key=_descriptor._internal_create_key),
+  ],
+  extensions=[
+  ],
+  nested_types=[],
+  enum_types=[
+    _LANECENTER_LANETYPE,
+  ],
+  serialized_options=None,
+  is_extendable=False,
+  syntax='proto2',
+  extension_ranges=[],
+  oneofs=[
+  ],
+  serialized_start=1385,
+  serialized_end=1892,
+)
+
+
+_ROADEDGE = _descriptor.Descriptor(
+  name='RoadEdge',
+  full_name='long_metric.RoadEdge',
+  filename=None,
+  file=DESCRIPTOR,
+  containing_type=None,
+  create_key=_descriptor._internal_create_key,
+  fields=[
+    _descriptor.FieldDescriptor(
+      name='type', full_name='long_metric.RoadEdge.type', index=0,
+      number=1, type=14, cpp_type=8, label=1,
+      has_default_value=False, default_value=0,
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR,  create_key=_descriptor._internal_create_key),
+    _descriptor.FieldDescriptor(
+      name='polyline', full_name='long_metric.RoadEdge.polyline', index=1,
+      number=2, type=11, cpp_type=10, label=3,
+      has_default_value=False, default_value=[],
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR,  create_key=_descriptor._internal_create_key),
+  ],
+  extensions=[
+  ],
+  nested_types=[],
+  enum_types=[
+    _ROADEDGE_ROADEDGETYPE,
+  ],
+  serialized_options=None,
+  is_extendable=False,
+  syntax='proto2',
+  extension_ranges=[],
+  oneofs=[
+  ],
+  serialized_start=1895,
+  serialized_end=2086,
+)
+
+
+_ROADLINE = _descriptor.Descriptor(
+  name='RoadLine',
+  full_name='long_metric.RoadLine',
+  filename=None,
+  file=DESCRIPTOR,
+  containing_type=None,
+  create_key=_descriptor._internal_create_key,
+  fields=[
+    _descriptor.FieldDescriptor(
+      name='type', full_name='long_metric.RoadLine.type', index=0,
+      number=1, type=14, cpp_type=8, label=1,
+      has_default_value=False, default_value=0,
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR,  create_key=_descriptor._internal_create_key),
+    _descriptor.FieldDescriptor(
+      name='polyline', full_name='long_metric.RoadLine.polyline', index=1,
+      number=2, type=11, cpp_type=10, label=3,
+      has_default_value=False, default_value=[],
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR,  create_key=_descriptor._internal_create_key),
+  ],
+  extensions=[
+  ],
+  nested_types=[],
+  enum_types=[
+    _ROADLINE_ROADLINETYPE,
+  ],
+  serialized_options=None,
+  is_extendable=False,
+  syntax='proto2',
+  extension_ranges=[],
+  oneofs=[
+  ],
+  serialized_start=2089,
+  serialized_end=2467,
+)
+
+
+_STOPSIGN = _descriptor.Descriptor(
+  name='StopSign',
+  full_name='long_metric.StopSign',
+  filename=None,
+  file=DESCRIPTOR,
+  containing_type=None,
+  create_key=_descriptor._internal_create_key,
+  fields=[
+    _descriptor.FieldDescriptor(
+      name='lane', full_name='long_metric.StopSign.lane', index=0,
+      number=1, type=3, cpp_type=2, label=3,
+      has_default_value=False, default_value=[],
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR,  create_key=_descriptor._internal_create_key),
+    _descriptor.FieldDescriptor(
+      name='position', full_name='long_metric.StopSign.position', index=1,
+      number=2, type=11, cpp_type=10, label=1,
+      has_default_value=False, default_value=None,
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR,  create_key=_descriptor._internal_create_key),
+  ],
+  extensions=[
+  ],
+  nested_types=[],
+  enum_types=[
+  ],
+  serialized_options=None,
+  is_extendable=False,
+  syntax='proto2',
+  extension_ranges=[],
+  oneofs=[
+  ],
+  serialized_start=2469,
+  serialized_end=2534,
+)
+
+
+_CROSSWALK = _descriptor.Descriptor(
+  name='Crosswalk',
+  full_name='long_metric.Crosswalk',
+  filename=None,
+  file=DESCRIPTOR,
+  containing_type=None,
+  create_key=_descriptor._internal_create_key,
+  fields=[
+    _descriptor.FieldDescriptor(
+      name='polygon', full_name='long_metric.Crosswalk.polygon', index=0,
+      number=1, type=11, cpp_type=10, label=3,
+      has_default_value=False, default_value=[],
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR,  create_key=_descriptor._internal_create_key),
+  ],
+  extensions=[
+  ],
+  nested_types=[],
+  enum_types=[
+  ],
+  serialized_options=None,
+  is_extendable=False,
+  syntax='proto2',
+  extension_ranges=[],
+  oneofs=[
+  ],
+  serialized_start=2536,
+  serialized_end=2587,
+)
+
+
+_SPEEDBUMP = _descriptor.Descriptor(
+  name='SpeedBump',
+  full_name='long_metric.SpeedBump',
+  filename=None,
+  file=DESCRIPTOR,
+  containing_type=None,
+  create_key=_descriptor._internal_create_key,
+  fields=[
+    _descriptor.FieldDescriptor(
+      name='polygon', full_name='long_metric.SpeedBump.polygon', index=0,
+      number=1, type=11, cpp_type=10, label=3,
+      has_default_value=False, default_value=[],
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR,  create_key=_descriptor._internal_create_key),
+  ],
+  extensions=[
+  ],
+  nested_types=[],
+  enum_types=[
+  ],
+  serialized_options=None,
+  is_extendable=False,
+  syntax='proto2',
+  extension_ranges=[],
+  oneofs=[
+  ],
+  serialized_start=2589,
+  serialized_end=2640,
+)
+
+
+_DRIVEWAY = _descriptor.Descriptor(
+  name='Driveway',
+  full_name='long_metric.Driveway',
+  filename=None,
+  file=DESCRIPTOR,
+  containing_type=None,
+  create_key=_descriptor._internal_create_key,
+  fields=[
+    _descriptor.FieldDescriptor(
+      name='polygon', full_name='long_metric.Driveway.polygon', index=0,
+      number=1, type=11, cpp_type=10, label=3,
+      has_default_value=False, default_value=[],
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR,  create_key=_descriptor._internal_create_key),
+  ],
+  extensions=[
+  ],
+  nested_types=[],
+  enum_types=[
+  ],
+  serialized_options=None,
+  is_extendable=False,
+  syntax='proto2',
+  extension_ranges=[],
+  oneofs=[
+  ],
+  serialized_start=2642,
+  serialized_end=2692,
+)
+
+_MAP.fields_by_name['map_features'].message_type = _MAPFEATURE
+_MAP.fields_by_name['dynamic_states'].message_type = _DYNAMICSTATE
+_DYNAMICSTATE.fields_by_name['lane_states'].message_type = _TRAFFICSIGNALLANESTATE
+_TRAFFICSIGNALLANESTATE.fields_by_name['state'].enum_type = _TRAFFICSIGNALLANESTATE_STATE
+_TRAFFICSIGNALLANESTATE.fields_by_name['stop_point'].message_type = _MAPPOINT
+_TRAFFICSIGNALLANESTATE_STATE.containing_type = _TRAFFICSIGNALLANESTATE
+_MAPFEATURE.fields_by_name['lane'].message_type = _LANECENTER
+_MAPFEATURE.fields_by_name['road_line'].message_type = _ROADLINE
+_MAPFEATURE.fields_by_name['road_edge'].message_type = _ROADEDGE
+_MAPFEATURE.fields_by_name['stop_sign'].message_type = _STOPSIGN
+_MAPFEATURE.fields_by_name['crosswalk'].message_type = _CROSSWALK
+_MAPFEATURE.fields_by_name['speed_bump'].message_type = _SPEEDBUMP
+_MAPFEATURE.fields_by_name['driveway'].message_type = _DRIVEWAY
+_MAPFEATURE.oneofs_by_name['feature_data'].fields.append(
+  _MAPFEATURE.fields_by_name['lane'])
+_MAPFEATURE.fields_by_name['lane'].containing_oneof = _MAPFEATURE.oneofs_by_name['feature_data']
+_MAPFEATURE.oneofs_by_name['feature_data'].fields.append(
+  _MAPFEATURE.fields_by_name['road_line'])
+_MAPFEATURE.fields_by_name['road_line'].containing_oneof = _MAPFEATURE.oneofs_by_name['feature_data']
+_MAPFEATURE.oneofs_by_name['feature_data'].fields.append(
+  _MAPFEATURE.fields_by_name['road_edge'])
+_MAPFEATURE.fields_by_name['road_edge'].containing_oneof = _MAPFEATURE.oneofs_by_name['feature_data']
+_MAPFEATURE.oneofs_by_name['feature_data'].fields.append(
+  _MAPFEATURE.fields_by_name['stop_sign'])
+_MAPFEATURE.fields_by_name['stop_sign'].containing_oneof = _MAPFEATURE.oneofs_by_name['feature_data']
+_MAPFEATURE.oneofs_by_name['feature_data'].fields.append(
+  _MAPFEATURE.fields_by_name['crosswalk'])
+_MAPFEATURE.fields_by_name['crosswalk'].containing_oneof = _MAPFEATURE.oneofs_by_name['feature_data']
+_MAPFEATURE.oneofs_by_name['feature_data'].fields.append(
+  _MAPFEATURE.fields_by_name['speed_bump'])
+_MAPFEATURE.fields_by_name['speed_bump'].containing_oneof = _MAPFEATURE.oneofs_by_name['feature_data']
+_MAPFEATURE.oneofs_by_name['feature_data'].fields.append(
+  _MAPFEATURE.fields_by_name['driveway'])
+_MAPFEATURE.fields_by_name['driveway'].containing_oneof = _MAPFEATURE.oneofs_by_name['feature_data']
+_BOUNDARYSEGMENT.fields_by_name['boundary_type'].enum_type = _ROADLINE_ROADLINETYPE
+_LANENEIGHBOR.fields_by_name['boundaries'].message_type = _BOUNDARYSEGMENT
+_LANECENTER.fields_by_name['type'].enum_type = _LANECENTER_LANETYPE
+_LANECENTER.fields_by_name['polyline'].message_type = _MAPPOINT
+_LANECENTER.fields_by_name['left_boundaries'].message_type = _BOUNDARYSEGMENT
+_LANECENTER.fields_by_name['right_boundaries'].message_type = _BOUNDARYSEGMENT
+_LANECENTER.fields_by_name['left_neighbors'].message_type = _LANENEIGHBOR
+_LANECENTER.fields_by_name['right_neighbors'].message_type = _LANENEIGHBOR
+_LANECENTER_LANETYPE.containing_type = _LANECENTER
+_ROADEDGE.fields_by_name['type'].enum_type = _ROADEDGE_ROADEDGETYPE
+_ROADEDGE.fields_by_name['polyline'].message_type = _MAPPOINT
+_ROADEDGE_ROADEDGETYPE.containing_type = _ROADEDGE
+_ROADLINE.fields_by_name['type'].enum_type = _ROADLINE_ROADLINETYPE
+_ROADLINE.fields_by_name['polyline'].message_type = _MAPPOINT
+_ROADLINE_ROADLINETYPE.containing_type = _ROADLINE
+_STOPSIGN.fields_by_name['position'].message_type = _MAPPOINT
+_CROSSWALK.fields_by_name['polygon'].message_type = _MAPPOINT
+_SPEEDBUMP.fields_by_name['polygon'].message_type = _MAPPOINT
+_DRIVEWAY.fields_by_name['polygon'].message_type = _MAPPOINT
+DESCRIPTOR.message_types_by_name['Map'] = _MAP
+DESCRIPTOR.message_types_by_name['DynamicState'] = _DYNAMICSTATE
+DESCRIPTOR.message_types_by_name['TrafficSignalLaneState'] = _TRAFFICSIGNALLANESTATE
+DESCRIPTOR.message_types_by_name['MapFeature'] = _MAPFEATURE
+DESCRIPTOR.message_types_by_name['MapPoint'] = _MAPPOINT
+DESCRIPTOR.message_types_by_name['BoundarySegment'] = _BOUNDARYSEGMENT
+DESCRIPTOR.message_types_by_name['LaneNeighbor'] = _LANENEIGHBOR
+DESCRIPTOR.message_types_by_name['LaneCenter'] = _LANECENTER
+DESCRIPTOR.message_types_by_name['RoadEdge'] = _ROADEDGE
+DESCRIPTOR.message_types_by_name['RoadLine'] = _ROADLINE
+DESCRIPTOR.message_types_by_name['StopSign'] = _STOPSIGN
+DESCRIPTOR.message_types_by_name['Crosswalk'] = _CROSSWALK
+DESCRIPTOR.message_types_by_name['SpeedBump'] = _SPEEDBUMP
+DESCRIPTOR.message_types_by_name['Driveway'] = _DRIVEWAY
+_sym_db.RegisterFileDescriptor(DESCRIPTOR)
+
+Map = _reflection.GeneratedProtocolMessageType('Map', (_message.Message,), {
+  'DESCRIPTOR' : _MAP,
+  '__module__' : 'dev.metrics.protos.map_pb2'
+  # @@protoc_insertion_point(class_scope:long_metric.Map)
+  })
+_sym_db.RegisterMessage(Map)
+
+DynamicState = _reflection.GeneratedProtocolMessageType('DynamicState', (_message.Message,), {
+  'DESCRIPTOR' : _DYNAMICSTATE,
+  '__module__' : 'dev.metrics.protos.map_pb2'
+  # @@protoc_insertion_point(class_scope:long_metric.DynamicState)
+  })
+_sym_db.RegisterMessage(DynamicState)
+
+TrafficSignalLaneState = _reflection.GeneratedProtocolMessageType('TrafficSignalLaneState', (_message.Message,), {
+  'DESCRIPTOR' : _TRAFFICSIGNALLANESTATE,
+  '__module__' : 'dev.metrics.protos.map_pb2'
+  # @@protoc_insertion_point(class_scope:long_metric.TrafficSignalLaneState)
+  })
+_sym_db.RegisterMessage(TrafficSignalLaneState)
+
+MapFeature = _reflection.GeneratedProtocolMessageType('MapFeature', (_message.Message,), {
+  'DESCRIPTOR' : _MAPFEATURE,
+  '__module__' : 'dev.metrics.protos.map_pb2'
+  # @@protoc_insertion_point(class_scope:long_metric.MapFeature)
+  })
+_sym_db.RegisterMessage(MapFeature)
+
+MapPoint = _reflection.GeneratedProtocolMessageType('MapPoint', (_message.Message,), {
+  'DESCRIPTOR' : _MAPPOINT,
+  '__module__' : 'dev.metrics.protos.map_pb2'
+  # @@protoc_insertion_point(class_scope:long_metric.MapPoint)
+  })
+_sym_db.RegisterMessage(MapPoint)
+
+BoundarySegment = _reflection.GeneratedProtocolMessageType('BoundarySegment', (_message.Message,), {
+  'DESCRIPTOR' : _BOUNDARYSEGMENT,
+  '__module__' : 'dev.metrics.protos.map_pb2'
+  # @@protoc_insertion_point(class_scope:long_metric.BoundarySegment)
+  })
+_sym_db.RegisterMessage(BoundarySegment)
+
+LaneNeighbor = _reflection.GeneratedProtocolMessageType('LaneNeighbor', (_message.Message,), {
+  'DESCRIPTOR' : _LANENEIGHBOR,
+  '__module__' : 'dev.metrics.protos.map_pb2'
+  # @@protoc_insertion_point(class_scope:long_metric.LaneNeighbor)
+  })
+_sym_db.RegisterMessage(LaneNeighbor)
+
+LaneCenter = _reflection.GeneratedProtocolMessageType('LaneCenter', (_message.Message,), {
+  'DESCRIPTOR' : _LANECENTER,
+  '__module__' : 'dev.metrics.protos.map_pb2'
+  # @@protoc_insertion_point(class_scope:long_metric.LaneCenter)
+  })
+_sym_db.RegisterMessage(LaneCenter)
+
+RoadEdge = _reflection.GeneratedProtocolMessageType('RoadEdge', (_message.Message,), {
+  'DESCRIPTOR' : _ROADEDGE,
+  '__module__' : 'dev.metrics.protos.map_pb2'
+  # @@protoc_insertion_point(class_scope:long_metric.RoadEdge)
+  })
+_sym_db.RegisterMessage(RoadEdge)
+
+RoadLine = _reflection.GeneratedProtocolMessageType('RoadLine', (_message.Message,), {
+  'DESCRIPTOR' : _ROADLINE,
+  '__module__' : 'dev.metrics.protos.map_pb2'
+  # @@protoc_insertion_point(class_scope:long_metric.RoadLine)
+  })
+_sym_db.RegisterMessage(RoadLine)
+
+StopSign = _reflection.GeneratedProtocolMessageType('StopSign', (_message.Message,), {
+  'DESCRIPTOR' : _STOPSIGN,
+  '__module__' : 'dev.metrics.protos.map_pb2'
+  # @@protoc_insertion_point(class_scope:long_metric.StopSign)
+  })
+_sym_db.RegisterMessage(StopSign)
+
+Crosswalk = _reflection.GeneratedProtocolMessageType('Crosswalk', (_message.Message,), {
+  'DESCRIPTOR' : _CROSSWALK,
+  '__module__' : 'dev.metrics.protos.map_pb2'
+  # @@protoc_insertion_point(class_scope:long_metric.Crosswalk)
+  })
+_sym_db.RegisterMessage(Crosswalk)
+
+SpeedBump = _reflection.GeneratedProtocolMessageType('SpeedBump', (_message.Message,), {
+  'DESCRIPTOR' : _SPEEDBUMP,
+  '__module__' : 'dev.metrics.protos.map_pb2'
+  # @@protoc_insertion_point(class_scope:long_metric.SpeedBump)
+  })
+_sym_db.RegisterMessage(SpeedBump)
+
+Driveway = _reflection.GeneratedProtocolMessageType('Driveway', (_message.Message,), {
+  'DESCRIPTOR' : _DRIVEWAY,
+  '__module__' : 'dev.metrics.protos.map_pb2'
+  # @@protoc_insertion_point(class_scope:long_metric.Driveway)
+  })
+_sym_db.RegisterMessage(Driveway)
+
+
+_LANECENTER.fields_by_name['entry_lanes']._options = None
+_LANECENTER.fields_by_name['exit_lanes']._options = None
+# @@protoc_insertion_point(module_scope)

seq_10_150_3_3_encode_occ_separate_offsets_bs8_128_no_seqindex_long_lastvalid_ablation_grid/backups/dev/metrics/protos/scenario_pb2.py

@@ -0,0 +1,454 @@
+# -*- coding: utf-8 -*-
+# Generated by the protocol buffer compiler.  DO NOT EDIT!
+# source: dev/metrics/protos/scenario.proto
+
+from google.protobuf import descriptor as _descriptor
+from google.protobuf import message as _message
+from google.protobuf import reflection as _reflection
+from google.protobuf import symbol_database as _symbol_database
+# @@protoc_insertion_point(imports)
+
+_sym_db = _symbol_database.Default()
+
+
+from dev.metrics.protos import map_pb2 as dev_dot_metrics_dot_protos_dot_map__pb2
+
+
+DESCRIPTOR = _descriptor.FileDescriptor(
+  name='dev/metrics/protos/scenario.proto',
+  package='long_metric',
+  syntax='proto2',
+  serialized_options=None,
+  create_key=_descriptor._internal_create_key,
+  serialized_pb=b'\n!dev/metrics/protos/scenario.proto\x12\x0blong_metric\x1a\x1c\x64\x65v/metrics/protos/map.proto\"\xba\x01\n\x0bObjectState\x12\x10\n\x08\x63\x65nter_x\x18\x02 \x01(\x01\x12\x10\n\x08\x63\x65nter_y\x18\x03 \x01(\x01\x12\x10\n\x08\x63\x65nter_z\x18\x04 \x01(\x01\x12\x0e\n\x06length\x18\x05 \x01(\x02\x12\r\n\x05width\x18\x06 \x01(\x02\x12\x0e\n\x06height\x18\x07 \x01(\x02\x12\x0f\n\x07heading\x18\x08 \x01(\x02\x12\x12\n\nvelocity_x\x18\t \x01(\x02\x12\x12\n\nvelocity_y\x18\n \x01(\x02\x12\r\n\x05valid\x18\x0b \x01(\x08\"\xd8\x01\n\x05Track\x12\n\n\x02id\x18\x01 \x01(\x05\x12\x32\n\x0bobject_type\x18\x02 \x01(\x0e\x32\x1d.long_metric.Track.ObjectType\x12(\n\x06states\x18\x03 \x03(\x0b\x32\x18.long_metric.ObjectState\"e\n\nObjectType\x12\x0e\n\nTYPE_UNSET\x10\x00\x12\x10\n\x0cTYPE_VEHICLE\x10\x01\x12\x13\n\x0fTYPE_PEDESTRIAN\x10\x02\x12\x10\n\x0cTYPE_CYCLIST\x10\x03\x12\x0e\n\nTYPE_OTHER\x10\x04\"K\n\x0f\x44ynamicMapState\x12\x38\n\x0blane_states\x18\x01 \x03(\x0b\x32#.long_metric.TrafficSignalLaneState\"\xa5\x01\n\x12RequiredPrediction\x12\x13\n\x0btrack_index\x18\x01 \x01(\x05\x12\x43\n\ndifficulty\x18\x02 \x01(\x0e\x32/.long_metric.RequiredPrediction.DifficultyLevel\"5\n\x0f\x44ifficultyLevel\x12\x08\n\x04NONE\x10\x00\x12\x0b\n\x07LEVEL_1\x10\x01\x12\x0b\n\x07LEVEL_2\x10\x02\"\xdc\x02\n\x08Scenario\x12\x13\n\x0bscenario_id\x18\x05 \x01(\t\x12\x1a\n\x12timestamps_seconds\x18\x01 \x03(\x01\x12\x1a\n\x12\x63urrent_time_index\x18\n \x01(\x05\x12\"\n\x06tracks\x18\x02 \x03(\x0b\x32\x12.long_metric.Track\x12\x38\n\x12\x64ynamic_map_states\x18\x07 \x03(\x0b\x32\x1c.long_metric.DynamicMapState\x12-\n\x0cmap_features\x18\x08 \x03(\x0b\x32\x17.long_metric.MapFeature\x12\x17\n\x0fsdc_track_index\x18\x06 \x01(\x05\x12\x1b\n\x13objects_of_interest\x18\x04 \x03(\x05\x12:\n\x11tracks_to_predict\x18\x0b \x03(\x0b\x32\x1f.long_metric.RequiredPredictionJ\x04\x08\t\x10\n'
+  ,
+  dependencies=[dev_dot_metrics_dot_protos_dot_map__pb2.DESCRIPTOR,])
+
+
+
+_TRACK_OBJECTTYPE = _descriptor.EnumDescriptor(
+  name='ObjectType',
+  full_name='long_metric.Track.ObjectType',
+  filename=None,
+  file=DESCRIPTOR,
+  create_key=_descriptor._internal_create_key,
+  values=[
+    _descriptor.EnumValueDescriptor(
+      name='TYPE_UNSET', index=0, number=0,
+      serialized_options=None,
+      type=None,
+      create_key=_descriptor._internal_create_key),
+    _descriptor.EnumValueDescriptor(
+      name='TYPE_VEHICLE', index=1, number=1,
+      serialized_options=None,
+      type=None,
+      create_key=_descriptor._internal_create_key),
+    _descriptor.EnumValueDescriptor(
+      name='TYPE_PEDESTRIAN', index=2, number=2,
+      serialized_options=None,
+      type=None,
+      create_key=_descriptor._internal_create_key),
+    _descriptor.EnumValueDescriptor(
+      name='TYPE_CYCLIST', index=3, number=3,
+      serialized_options=None,
+      type=None,
+      create_key=_descriptor._internal_create_key),
+    _descriptor.EnumValueDescriptor(
+      name='TYPE_OTHER', index=4, number=4,
+      serialized_options=None,
+      type=None,
+      create_key=_descriptor._internal_create_key),
+  ],
+  containing_type=None,
+  serialized_options=None,
+  serialized_start=385,
+  serialized_end=486,
+)
+_sym_db.RegisterEnumDescriptor(_TRACK_OBJECTTYPE)
+
+_REQUIREDPREDICTION_DIFFICULTYLEVEL = _descriptor.EnumDescriptor(
+  name='DifficultyLevel',
+  full_name='long_metric.RequiredPrediction.DifficultyLevel',
+  filename=None,
+  file=DESCRIPTOR,
+  create_key=_descriptor._internal_create_key,
+  values=[
+    _descriptor.EnumValueDescriptor(
+      name='NONE', index=0, number=0,
+      serialized_options=None,
+      type=None,
+      create_key=_descriptor._internal_create_key),
+    _descriptor.EnumValueDescriptor(
+      name='LEVEL_1', index=1, number=1,
+      serialized_options=None,
+      type=None,
+      create_key=_descriptor._internal_create_key),
+    _descriptor.EnumValueDescriptor(
+      name='LEVEL_2', index=2, number=2,
+      serialized_options=None,
+      type=None,
+      create_key=_descriptor._internal_create_key),
+  ],
+  containing_type=None,
+  serialized_options=None,
+  serialized_start=678,
+  serialized_end=731,
+)
+_sym_db.RegisterEnumDescriptor(_REQUIREDPREDICTION_DIFFICULTYLEVEL)
+
+
+_OBJECTSTATE = _descriptor.Descriptor(
+  name='ObjectState',
+  full_name='long_metric.ObjectState',
+  filename=None,
+  file=DESCRIPTOR,
+  containing_type=None,
+  create_key=_descriptor._internal_create_key,
+  fields=[
+    _descriptor.FieldDescriptor(
+      name='center_x', full_name='long_metric.ObjectState.center_x', index=0,
+      number=2, type=1, cpp_type=5, label=1,
+      has_default_value=False, default_value=float(0),
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR,  create_key=_descriptor._internal_create_key),
+    _descriptor.FieldDescriptor(
+      name='center_y', full_name='long_metric.ObjectState.center_y', index=1,
+      number=3, type=1, cpp_type=5, label=1,
+      has_default_value=False, default_value=float(0),
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR,  create_key=_descriptor._internal_create_key),
+    _descriptor.FieldDescriptor(
+      name='center_z', full_name='long_metric.ObjectState.center_z', index=2,
+      number=4, type=1, cpp_type=5, label=1,
+      has_default_value=False, default_value=float(0),
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR,  create_key=_descriptor._internal_create_key),
+    _descriptor.FieldDescriptor(
+      name='length', full_name='long_metric.ObjectState.length', index=3,
+      number=5, type=2, cpp_type=6, label=1,
+      has_default_value=False, default_value=float(0),
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR,  create_key=_descriptor._internal_create_key),
+    _descriptor.FieldDescriptor(
+      name='width', full_name='long_metric.ObjectState.width', index=4,
+      number=6, type=2, cpp_type=6, label=1,
+      has_default_value=False, default_value=float(0),
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR,  create_key=_descriptor._internal_create_key),
+    _descriptor.FieldDescriptor(
+      name='height', full_name='long_metric.ObjectState.height', index=5,
+      number=7, type=2, cpp_type=6, label=1,
+      has_default_value=False, default_value=float(0),
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR,  create_key=_descriptor._internal_create_key),
+    _descriptor.FieldDescriptor(
+      name='heading', full_name='long_metric.ObjectState.heading', index=6,
+      number=8, type=2, cpp_type=6, label=1,
+      has_default_value=False, default_value=float(0),
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR,  create_key=_descriptor._internal_create_key),
+    _descriptor.FieldDescriptor(
+      name='velocity_x', full_name='long_metric.ObjectState.velocity_x', index=7,
+      number=9, type=2, cpp_type=6, label=1,
+      has_default_value=False, default_value=float(0),
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR,  create_key=_descriptor._internal_create_key),
+    _descriptor.FieldDescriptor(
+      name='velocity_y', full_name='long_metric.ObjectState.velocity_y', index=8,
+      number=10, type=2, cpp_type=6, label=1,
+      has_default_value=False, default_value=float(0),
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR,  create_key=_descriptor._internal_create_key),
+    _descriptor.FieldDescriptor(
+      name='valid', full_name='long_metric.ObjectState.valid', index=9,
+      number=11, type=8, cpp_type=7, label=1,
+      has_default_value=False, default_value=False,
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR,  create_key=_descriptor._internal_create_key),
+  ],
+  extensions=[
+  ],
+  nested_types=[],
+  enum_types=[
+  ],
+  serialized_options=None,
+  is_extendable=False,
+  syntax='proto2',
+  extension_ranges=[],
+  oneofs=[
+  ],
+  serialized_start=81,
+  serialized_end=267,
+)
+
+
+_TRACK = _descriptor.Descriptor(
+  name='Track',
+  full_name='long_metric.Track',
+  filename=None,
+  file=DESCRIPTOR,
+  containing_type=None,
+  create_key=_descriptor._internal_create_key,
+  fields=[
+    _descriptor.FieldDescriptor(
+      name='id', full_name='long_metric.Track.id', index=0,
+      number=1, type=5, cpp_type=1, label=1,
+      has_default_value=False, default_value=0,
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR,  create_key=_descriptor._internal_create_key),
+    _descriptor.FieldDescriptor(
+      name='object_type', full_name='long_metric.Track.object_type', index=1,
+      number=2, type=14, cpp_type=8, label=1,
+      has_default_value=False, default_value=0,
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR,  create_key=_descriptor._internal_create_key),
+    _descriptor.FieldDescriptor(
+      name='states', full_name='long_metric.Track.states', index=2,
+      number=3, type=11, cpp_type=10, label=3,
+      has_default_value=False, default_value=[],
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR,  create_key=_descriptor._internal_create_key),
+  ],
+  extensions=[
+  ],
+  nested_types=[],
+  enum_types=[
+    _TRACK_OBJECTTYPE,
+  ],
+  serialized_options=None,
+  is_extendable=False,
+  syntax='proto2',
+  extension_ranges=[],
+  oneofs=[
+  ],
+  serialized_start=270,
+  serialized_end=486,
+)
+
+
+_DYNAMICMAPSTATE = _descriptor.Descriptor(
+  name='DynamicMapState',
+  full_name='long_metric.DynamicMapState',
+  filename=None,
+  file=DESCRIPTOR,
+  containing_type=None,
+  create_key=_descriptor._internal_create_key,
+  fields=[
+    _descriptor.FieldDescriptor(
+      name='lane_states', full_name='long_metric.DynamicMapState.lane_states', index=0,
+      number=1, type=11, cpp_type=10, label=3,
+      has_default_value=False, default_value=[],
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR,  create_key=_descriptor._internal_create_key),
+  ],
+  extensions=[
+  ],
+  nested_types=[],
+  enum_types=[
+  ],
+  serialized_options=None,
+  is_extendable=False,
+  syntax='proto2',
+  extension_ranges=[],
+  oneofs=[
+  ],
+  serialized_start=488,
+  serialized_end=563,
+)
+
+
+_REQUIREDPREDICTION = _descriptor.Descriptor(
+  name='RequiredPrediction',
+  full_name='long_metric.RequiredPrediction',
+  filename=None,
+  file=DESCRIPTOR,
+  containing_type=None,
+  create_key=_descriptor._internal_create_key,
+  fields=[
+    _descriptor.FieldDescriptor(
+      name='track_index', full_name='long_metric.RequiredPrediction.track_index', index=0,
+      number=1, type=5, cpp_type=1, label=1,
+      has_default_value=False, default_value=0,
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR,  create_key=_descriptor._internal_create_key),
+    _descriptor.FieldDescriptor(
+      name='difficulty', full_name='long_metric.RequiredPrediction.difficulty', index=1,
+      number=2, type=14, cpp_type=8, label=1,
+      has_default_value=False, default_value=0,
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR,  create_key=_descriptor._internal_create_key),
+  ],
+  extensions=[
+  ],
+  nested_types=[],
+  enum_types=[
+    _REQUIREDPREDICTION_DIFFICULTYLEVEL,
+  ],
+  serialized_options=None,
+  is_extendable=False,
+  syntax='proto2',
+  extension_ranges=[],
+  oneofs=[
+  ],
+  serialized_start=566,
+  serialized_end=731,
+)
+
+
+_SCENARIO = _descriptor.Descriptor(
+  name='Scenario',
+  full_name='long_metric.Scenario',
+  filename=None,
+  file=DESCRIPTOR,
+  containing_type=None,
+  create_key=_descriptor._internal_create_key,
+  fields=[
+    _descriptor.FieldDescriptor(
+      name='scenario_id', full_name='long_metric.Scenario.scenario_id', index=0,
+      number=5, type=9, cpp_type=9, label=1,
+      has_default_value=False, default_value=b"".decode('utf-8'),
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR,  create_key=_descriptor._internal_create_key),
+    _descriptor.FieldDescriptor(
+      name='timestamps_seconds', full_name='long_metric.Scenario.timestamps_seconds', index=1,
+      number=1, type=1, cpp_type=5, label=3,
+      has_default_value=False, default_value=[],
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR,  create_key=_descriptor._internal_create_key),
+    _descriptor.FieldDescriptor(
+      name='current_time_index', full_name='long_metric.Scenario.current_time_index', index=2,
+      number=10, type=5, cpp_type=1, label=1,
+      has_default_value=False, default_value=0,
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR,  create_key=_descriptor._internal_create_key),
+    _descriptor.FieldDescriptor(
+      name='tracks', full_name='long_metric.Scenario.tracks', index=3,
+      number=2, type=11, cpp_type=10, label=3,
+      has_default_value=False, default_value=[],
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR,  create_key=_descriptor._internal_create_key),
+    _descriptor.FieldDescriptor(
+      name='dynamic_map_states', full_name='long_metric.Scenario.dynamic_map_states', index=4,
+      number=7, type=11, cpp_type=10, label=3,
+      has_default_value=False, default_value=[],
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR,  create_key=_descriptor._internal_create_key),
+    _descriptor.FieldDescriptor(
+      name='map_features', full_name='long_metric.Scenario.map_features', index=5,
+      number=8, type=11, cpp_type=10, label=3,
+      has_default_value=False, default_value=[],
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR,  create_key=_descriptor._internal_create_key),
+    _descriptor.FieldDescriptor(
+      name='sdc_track_index', full_name='long_metric.Scenario.sdc_track_index', index=6,
+      number=6, type=5, cpp_type=1, label=1,
+      has_default_value=False, default_value=0,
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR,  create_key=_descriptor._internal_create_key),
+    _descriptor.FieldDescriptor(
+      name='objects_of_interest', full_name='long_metric.Scenario.objects_of_interest', index=7,
+      number=4, type=5, cpp_type=1, label=3,
+      has_default_value=False, default_value=[],
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR,  create_key=_descriptor._internal_create_key),
+    _descriptor.FieldDescriptor(
+      name='tracks_to_predict', full_name='long_metric.Scenario.tracks_to_predict', index=8,
+      number=11, type=11, cpp_type=10, label=3,
+      has_default_value=False, default_value=[],
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR,  create_key=_descriptor._internal_create_key),
+  ],
+  extensions=[
+  ],
+  nested_types=[],
+  enum_types=[
+  ],
+  serialized_options=None,
+  is_extendable=False,
+  syntax='proto2',
+  extension_ranges=[],
+  oneofs=[
+  ],
+  serialized_start=734,
+  serialized_end=1082,
+)
+
+_TRACK.fields_by_name['object_type'].enum_type = _TRACK_OBJECTTYPE
+_TRACK.fields_by_name['states'].message_type = _OBJECTSTATE
+_TRACK_OBJECTTYPE.containing_type = _TRACK
+_DYNAMICMAPSTATE.fields_by_name['lane_states'].message_type = dev_dot_metrics_dot_protos_dot_map__pb2._TRAFFICSIGNALLANESTATE
+_REQUIREDPREDICTION.fields_by_name['difficulty'].enum_type = _REQUIREDPREDICTION_DIFFICULTYLEVEL
+_REQUIREDPREDICTION_DIFFICULTYLEVEL.containing_type = _REQUIREDPREDICTION
+_SCENARIO.fields_by_name['tracks'].message_type = _TRACK
+_SCENARIO.fields_by_name['dynamic_map_states'].message_type = _DYNAMICMAPSTATE
+_SCENARIO.fields_by_name['map_features'].message_type = dev_dot_metrics_dot_protos_dot_map__pb2._MAPFEATURE
+_SCENARIO.fields_by_name['tracks_to_predict'].message_type = _REQUIREDPREDICTION
+DESCRIPTOR.message_types_by_name['ObjectState'] = _OBJECTSTATE
+DESCRIPTOR.message_types_by_name['Track'] = _TRACK
+DESCRIPTOR.message_types_by_name['DynamicMapState'] = _DYNAMICMAPSTATE
+DESCRIPTOR.message_types_by_name['RequiredPrediction'] = _REQUIREDPREDICTION
+DESCRIPTOR.message_types_by_name['Scenario'] = _SCENARIO
+_sym_db.RegisterFileDescriptor(DESCRIPTOR)
+
+ObjectState = _reflection.GeneratedProtocolMessageType('ObjectState', (_message.Message,), {
+  'DESCRIPTOR' : _OBJECTSTATE,
+  '__module__' : 'dev.metrics.protos.scenario_pb2'
+  # @@protoc_insertion_point(class_scope:long_metric.ObjectState)
+  })
+_sym_db.RegisterMessage(ObjectState)
+
+Track = _reflection.GeneratedProtocolMessageType('Track', (_message.Message,), {
+  'DESCRIPTOR' : _TRACK,
+  '__module__' : 'dev.metrics.protos.scenario_pb2'
+  # @@protoc_insertion_point(class_scope:long_metric.Track)
+  })
+_sym_db.RegisterMessage(Track)
+
+DynamicMapState = _reflection.GeneratedProtocolMessageType('DynamicMapState', (_message.Message,), {
+  'DESCRIPTOR' : _DYNAMICMAPSTATE,
+  '__module__' : 'dev.metrics.protos.scenario_pb2'
+  # @@protoc_insertion_point(class_scope:long_metric.DynamicMapState)
+  })
+_sym_db.RegisterMessage(DynamicMapState)
+
+RequiredPrediction = _reflection.GeneratedProtocolMessageType('RequiredPrediction', (_message.Message,), {
+  'DESCRIPTOR' : _REQUIREDPREDICTION,
+  '__module__' : 'dev.metrics.protos.scenario_pb2'
+  # @@protoc_insertion_point(class_scope:long_metric.RequiredPrediction)
+  })
+_sym_db.RegisterMessage(RequiredPrediction)
+
+Scenario = _reflection.GeneratedProtocolMessageType('Scenario', (_message.Message,), {
+  'DESCRIPTOR' : _SCENARIO,
+  '__module__' : 'dev.metrics.protos.scenario_pb2'
+  # @@protoc_insertion_point(class_scope:long_metric.Scenario)
+  })
+_sym_db.RegisterMessage(Scenario)
+
+
+# @@protoc_insertion_point(module_scope)

seq_10_150_3_3_encode_occ_separate_offsets_bs8_128_no_seqindex_long_lastvalid_ablation_grid/backups/dev/metrics/trajectory_features.py

@@ -0,0 +1,52 @@
+import torch
+import numpy as np
+from torch import Tensor
+from typing import Tuple
+
+
+def _wrap_angle(angle: Tensor) -> Tensor:
+    return (angle + np.pi) % (2 * np.pi) - np.pi
+
+
+def central_diff(t: Tensor, pad_value: float) -> Tensor:
+    pad_shape = (*t.shape[:-1], 1)
+    pad_tensor = torch.full(pad_shape, pad_value, dtype=t.dtype, device=t.device)
+    diff_t = (t[..., 2:] - t[..., :-2]) / 2
+    return torch.cat([pad_tensor, diff_t, pad_tensor], dim=-1)
+
+
+def central_logical_and(t: Tensor, pad_value: bool) -> Tensor:
+    pad_shape = (*t.shape[:-1], 1)
+    pad_tensor = torch.full(pad_shape, pad_value, dtype=torch.bool, device=t.device)
+    diff_t = torch.logical_and(t[..., 2:], t[..., :-2])
+    return torch.cat([pad_tensor, diff_t, pad_tensor], dim=-1)
+
+
+def compute_displacement_error(x, y, z, ref_x, ref_y, ref_z) -> Tensor:
+    return torch.norm(
+        torch.stack([x, y, z], dim=-1) - torch.stack([ref_x, ref_y, ref_z], dim=-1),
+        p=2, dim=-1
+    )
+
+
+def compute_kinematic_features(
+    x: Tensor,
+    y: Tensor,
+    z: Tensor,
+    heading: Tensor,
+    seconds_per_step: float
+) -> Tuple[Tensor, Tensor, Tensor, Tensor]:
+    dpos = central_diff(torch.stack([x, y, z], dim=0), pad_value=np.nan)
+    linear_speed = torch.norm(dpos, p=2, dim=0) / seconds_per_step
+    linear_accel = central_diff(linear_speed, pad_value=np.nan) / seconds_per_step
+    dh_step = _wrap_angle(central_diff(heading, pad_value=np.nan) * 2) / 2
+    dh = dh_step / seconds_per_step
+    d2h_step = _wrap_angle(central_diff(dh_step, pad_value=np.nan) * 2) / 2
+    d2h = d2h_step / (seconds_per_step ** 2)
+    return linear_speed, linear_accel, dh, d2h
+
+
+def compute_kinematic_validity(valid: Tensor) -> Tuple[Tensor, Tensor]:
+    speed_validity = central_logical_and(valid, pad_value=False)
+    acceleration_validity = central_logical_and(speed_validity, pad_value=False)
+    return speed_validity, acceleration_validity

seq_10_150_3_3_encode_occ_separate_offsets_bs8_128_no_seqindex_long_lastvalid_ablation_grid/backups/dev/metrics/val_close_long_metrics.json

@@ -0,0 +1,24 @@
+{
+    "val_close_long/wosac/realism_meta_metric": 0.6187323331832886,
+    "val_close_long/wosac/kinematic_metrics": 0.6323384046554565,
+    "val_close_long/wosac/interactive_metrics": 0.5528579354286194,
+    "val_close_long/wosac/map_based_metrics": 0.0,
+    "val_close_long/wosac/placement_based_metrics": 0.6086956858634949,
+    "val_close_long/wosac/min_ade": 0.0,
+    "val_close_long/wosac/scenario_counter": 61,
+    "val_close_long/wosac_likelihood/metametric": 0.6187323331832886,
+    "val_close_long/wosac_likelihood/average_displacement_error": 0.0,
+    "val_close_long/wosac_likelihood/min_average_displacement_error": 0.0,
+    "val_close_long/wosac_likelihood/linear_speed_likelihood": 0.11858943104743958,
+    "val_close_long/wosac_likelihood/linear_acceleration_likelihood": 0.6093839406967163,
+    "val_close_long/wosac_likelihood/angular_speed_likelihood": 0.8988037705421448,
+    "val_close_long/wosac_likelihood/angular_acceleration_likelihood": 0.9025763869285583,
+    "val_close_long/wosac_likelihood/distance_to_nearest_object_likelihood": 0.10390616208314896,
+    "val_close_long/wosac_likelihood/collision_indication_likelihood": 0.6108496785163879,
+    "val_close_long/wosac_likelihood/time_to_collision_likelihood": 0.8568302989006042,
+    "val_close_long/wosac_likelihood/simulated_collision_rate": 0.030917881056666374,
+    "val_close_long/wosac_likelihood/num_placement_likelihood": 0.7245867848396301,
+    "val_close_long/wosac_likelihood/num_removement_likelihood": 0.6228984594345093,
+    "val_close_long/wosac_likelihood/distance_placement_likelihood": 1.0,
+    "val_close_long/wosac_likelihood/distance_removement_likelihood": 0.08729743212461472
+}

seq_10_150_3_3_encode_occ_separate_offsets_bs8_128_no_seqindex_long_lastvalid_ablation_grid/backups/dev/model/smart.py

@@ -0,0 +1,1100 @@
+import os
+import contextlib
+import pytorch_lightning as pl
+import math
+import numpy as np
+import pickle
+import random
+import torch
+import torch.nn as nn
+from tqdm import tqdm
+from torch_geometric.data import Batch
+from torch_geometric.data import HeteroData
+from torch.optim.lr_scheduler import LambdaLR
+from collections import defaultdict
+
+from dev.utils.func import angle_between_2d_vectors
+from dev.modules.layers import OccLoss
+from dev.modules.attr_tokenizer import Attr_Tokenizer
+from dev.modules.smart_decoder import SMARTDecoder
+from dev.datasets.preprocess import TokenProcessor
+from dev.metrics.compute_metrics import *
+from dev.utils.metrics import *
+from dev.utils.visualization import *
+
+
+class SMART(pl.LightningModule):
+
+    def __init__(self, model_config, save_path: os.PathLike="", logger=None, **kwargs) -> None:
+        super(SMART, self).__init__()
+        self.save_hyperparameters()
+        self.model_config = model_config
+        self.warmup_steps = model_config.warmup_steps
+        self.lr = model_config.lr
+        self.total_steps = model_config.total_steps
+        self.dataset = model_config.dataset
+        self.input_dim = model_config.input_dim
+        self.hidden_dim = model_config.hidden_dim
+        self.output_dim = model_config.output_dim
+        self.output_head = model_config.output_head
+        self.num_historical_steps = model_config.num_historical_steps
+        self.num_future_steps = model_config.decoder.num_future_steps
+        self.num_freq_bands = model_config.num_freq_bands
+        self.save_path = save_path
+        self.vis_map = False
+        self.noise = True
+        self.local_logger = logger
+        self.max_epochs = kwargs.get('max_epochs', 0)
+        module_dir = os.path.dirname(os.path.dirname(__file__))
+
+        self.map_token_traj_path = os.path.join(module_dir, 'tokens/map_traj_token5.pkl')
+        self.init_map_token()
+
+        self.predict_motion = model_config.predict_motion
+        self.predict_state = model_config.predict_state
+        self.predict_map = model_config.predict_map
+        self.predict_occ = model_config.predict_occ
+        self.pl2seed_radius = model_config.decoder.pl2seed_radius
+        self.token_size = model_config.decoder.token_size
+
+        # if `disable_grid_token` is True, then we process all locations as
+        # the continuous values. Besides, no occupancy grid input.
+        # Also, no need to predict the xy offset.
+        self.disable_grid_token = getattr(model_config, 'disable_grid_token') \
+            if hasattr(model_config, 'disable_grid_token') else False
+        self.use_grid_token = not self.disable_grid_token
+        if self.disable_grid_token:
+            self.predict_occ = False
+
+        self.token_processer = TokenProcessor(self.token_size,
+                                              training=self.training,
+                                              predict_motion=self.predict_motion,
+                                              predict_state=self.predict_state,
+                                              predict_map=self.predict_map,
+                                              state_token=model_config.state_token,
+                                              pl2seed_radius=self.pl2seed_radius)
+
+        self.attr_tokenizer = Attr_Tokenizer(grid_range=self.model_config.grid_range,
+                                             grid_interval=self.model_config.grid_interval,
+                                             radius=model_config.decoder.pl2seed_radius,
+                                             angle_interval=self.model_config.angle_interval)
+
+        # state tokens
+        self.invalid_state = int(self.model_config.state_token['invalid'])
+        self.valid_state = int(self.model_config.state_token['valid'])
+        self.enter_state = int(self.model_config.state_token['enter'])
+        self.exit_state = int(self.model_config.state_token['exit'])
+
+        self.seed_size = int(model_config.decoder.seed_size)
+
+        self.encoder = SMARTDecoder(
+            decoder_type=model_config.decoder_type,
+            dataset=model_config.dataset,
+            input_dim=model_config.input_dim,
+            hidden_dim=model_config.hidden_dim,
+            num_historical_steps=model_config.num_historical_steps,
+            num_freq_bands=model_config.num_freq_bands,
+            num_heads=model_config.num_heads,
+            head_dim=model_config.head_dim,
+            dropout=model_config.dropout,
+            num_map_layers=model_config.decoder.num_map_layers,
+            num_agent_layers=model_config.decoder.num_agent_layers,
+            pl2pl_radius=model_config.decoder.pl2pl_radius,
+            pl2a_radius=model_config.decoder.pl2a_radius,
+            pl2seed_radius=model_config.decoder.pl2seed_radius,
+            a2a_radius=model_config.decoder.a2a_radius,
+            a2sa_radius=model_config.decoder.a2sa_radius,
+            pl2sa_radius=model_config.decoder.pl2sa_radius,
+            time_span=model_config.decoder.time_span,
+            map_token={'traj_src': self.map_token['traj_src']},
+            token_size=self.token_size,
+            attr_tokenizer=self.attr_tokenizer,
+            predict_motion=self.predict_motion,
+            predict_state=self.predict_state,
+            predict_map=self.predict_map,
+            predict_occ=self.predict_occ,
+            state_token=model_config.state_token,
+            use_grid_token=self.use_grid_token,
+            seed_size=self.seed_size,
+            buffer_size=model_config.decoder.buffer_size,
+            num_recurrent_steps_val=model_config.num_recurrent_steps_val,
+            loss_weight=model_config.loss_weight,
+            logger=logger,
+        )
+        self.minADE = minADE(max_guesses=1)
+        self.minFDE = minFDE(max_guesses=1)
+        self.TokenCls = TokenCls(max_guesses=1)
+        self.StateCls = TokenCls(max_guesses=1)
+        self.StateAccuracy = StateAccuracy(state_token=self.model_config.state_token)
+        self.GridOverlapRate = GridOverlapRate(num_step=18,
+                                               state_token=self.model_config.state_token,
+                                               seed_size=self.encoder.agent_encoder.num_seed_feature)
+        # self.NumInsertAccuracy = NumInsertAccuracy()
+        self.loss_weight = model_config.loss_weight
+
+        self.token_cls_loss = nn.CrossEntropyLoss(label_smoothing=0.1)
+        if self.predict_map:
+            self.map_token_loss = nn.CrossEntropyLoss(label_smoothing=0.1)
+        if self.predict_state:
+            self.state_cls_loss = nn.CrossEntropyLoss(
+                torch.tensor(self.loss_weight['state_weight']))
+            self.state_cls_loss_seed = nn.CrossEntropyLoss(
+                torch.tensor(self.loss_weight['seed_state_weight']))
+            self.type_cls_loss_seed = nn.CrossEntropyLoss(
+                torch.tensor(self.loss_weight['seed_type_weight']))
+            self.pos_cls_loss_seed = nn.CrossEntropyLoss(label_smoothing=0.1)
+            self.head_cls_loss_seed = nn.CrossEntropyLoss()
+            self.offset_reg_loss_seed = nn.MSELoss()
+            self.shape_reg_loss_seed = nn.MSELoss()
+            self.pos_reg_loss_seed = nn.MSELoss()
+        if self.predict_occ:
+            self.occ_cls_loss = nn.CrossEntropyLoss()
+            self.agent_occ_loss_seed = nn.BCEWithLogitsLoss(
+                pos_weight=torch.tensor([self.loss_weight['agent_occ_pos_weight']]))
+            self.pt_occ_loss_seed = nn.BCEWithLogitsLoss(
+                pos_weight=torch.tensor([self.loss_weight['pt_occ_pos_weight']]))
+            # self.agent_occ_loss_seed = OccLoss()
+            # self.pt_occ_loss_seed = OccLoss()
+            # self.agent_occ_loss_seed = nn.BCEWithLogitsLoss()
+            # self.pt_occ_loss_seed = nn.BCEWithLogitsLoss()
+        self.rollout_num = 1
+
+        self.val_open_loop = model_config.val_open_loop
+        self.val_close_loop = model_config.val_close_loop
+        self.val_insert = model_config.val_insert or bool(os.getenv('VAL_INSERT'))
+        self.n_rollout_close_val = model_config.n_rollout_close_val
+        self.t = kwargs.get('t', 2)
+
+        # for validation / test
+        self._mode = 'training'
+        self._long_metrics = None
+        self._online_metric = False
+        self._save_validate_reuslts = False
+        self._plot_rollouts = False
+
+    def set(self, mode: str = 'train'):
+        self._mode = mode
+
+        if mode == 'validation':
+            self._online_metric = True
+            self._save_validate_reuslts = True
+            self._long_metrics = LongMetric('val_close_long')
+
+        elif mode == 'test':
+            self._save_validate_reuslts = True
+
+        elif mode == 'plot_rollouts':
+            self._plot_rollouts = True
+
+    def init_map_token(self):
+        self.argmin_sample_len = 3
+        map_token_traj = pickle.load(open(self.map_token_traj_path, 'rb'))
+        self.map_token = {'traj_src': map_token_traj['traj_src'], }
+        traj_end_theta = np.arctan2(self.map_token['traj_src'][:, -1, 1]-self.map_token['traj_src'][:, -2, 1],
+                                    self.map_token['traj_src'][:, -1, 0]-self.map_token['traj_src'][:, -2, 0])
+        indices = torch.linspace(0, self.map_token['traj_src'].shape[1]-1, steps=self.argmin_sample_len).long()
+        self.map_token['sample_pt'] = torch.from_numpy(self.map_token['traj_src'][:, indices]).to(torch.float)
+        self.map_token['traj_end_theta'] = torch.from_numpy(traj_end_theta).to(torch.float)
+        self.map_token['traj_src'] = torch.from_numpy(self.map_token['traj_src']).to(torch.float)
+
+    def get_agent_inputs(self, data: HeteroData):
+        res = self.encoder.get_agent_inputs(data)
+        return res
+
+    def forward(self, data: HeteroData):
+        res = self.encoder(data)
+        return res
+
+    def maybe_autocast(self, dtype=torch.float16):
+        enable_autocast = self.device != torch.device("cpu")
+
+        if enable_autocast:
+            return torch.cuda.amp.autocast(dtype=dtype)
+        else:
+            return contextlib.nullcontext()
+
+    def check_inputs(self, data: HeteroData):
+        inputs = self.get_agent_inputs(data)
+        next_token_idx_gt = inputs['next_token_idx_gt']
+        next_state_idx_gt = inputs['next_state_idx_gt'].clone()
+        next_token_eval_mask = inputs['next_token_eval_mask'].clone()
+        raw_agent_valid_mask = inputs['raw_agent_valid_mask'].clone()
+
+        self.StateAccuracy.update(state_idx=next_state_idx_gt,
+                                  valid_mask=raw_agent_valid_mask)
+
+        state_token = inputs['state_token']
+        grid_index = inputs['grid_index']
+        self.GridOverlapRate.update(state_token=state_token,
+                                    grid_index=grid_index)
+
+        print(self.StateAccuracy)
+        print(self.GridOverlapRate)
+        # self.log('valid_accuracy', self.StateAccuracy.compute()['valid'], prog_bar=True, on_step=True, on_epoch=True, batch_size=1)
+        # self.log('invalid_accuracy', self.StateAccuracy.compute()['invalid'], prog_bar=True, on_step=True, on_epoch=True, batch_size=1)
+
+    def training_step(self,
+                      data,
+                      batch_idx):
+
+        data = self.token_processer(data)
+
+        data = self.match_token_map(data)
+        data = self.sample_pt_pred(data)
+
+        # find map tokens for entering agents
+        data = self._fetch_enterings(data)
+
+        data['batch_size_a'] = data['agent']['ptr'][1:] - data['agent']['ptr'][:-1]
+        data['batch_size_pl'] = data['pt_token']['ptr'][1:] - data['pt_token']['ptr'][:-1]
+        if isinstance(data, Batch):
+            data['agent']['av_index'] += data['agent']['ptr'][:-1]
+
+        if int(os.getenv("CHECK_INPUTS", 0)):
+            return self.check_inputs(data)
+
+        pred = self(data)
+
+        loss = 0
+
+        log_params = dict(prog_bar=True, on_step=True, on_epoch=True, batch_size=1, sync_dist=True)
+
+        if pred.get('occ_decoder', False):
+
+            agent_occ = pred['agent_occ']
+            agent_occ_gt = pred['agent_occ_gt']
+            agent_occ_eval_mask = pred['agent_occ_eval_mask']
+            pt_occ = pred['pt_occ']
+            pt_occ_gt = pred['pt_occ_gt']
+            pt_occ_eval_mask = pred['pt_occ_eval_mask']
+
+            agent_occ_cls_loss = self.occ_cls_loss(agent_occ[agent_occ_eval_mask],
+                                                   agent_occ_gt[agent_occ_eval_mask])
+            pt_occ_cls_loss = self.occ_cls_loss(pt_occ[pt_occ_eval_mask],
+                                                pt_occ_gt[pt_occ_eval_mask])
+            self.log('agent_occ_cls_loss', agent_occ_cls_loss, **log_params)
+            self.log('pt_occ_cls_loss', pt_occ_cls_loss, **log_params)
+            loss = loss + agent_occ_cls_loss + pt_occ_cls_loss
+
+            # plot
+            # plot_scenario_ids = ['74ad7b76d5906d39', '1351ea8b8333ddcb', '1352066cc3c0508d', '135436833ce5b9e7', '13570a32432d449', '13577c32a81336fb']
+            if random.random() < 4e-5 or os.getenv('DEBUG'):
+                num_step = pred['num_step']
+                num_agent = pred['num_agent']
+                num_pt = pred['num_pt']
+                with torch.no_grad():
+                    agent_occ = agent_occ.reshape(num_step, num_agent, -1).transpose(0, 1)
+                    agent_occ_gt = agent_occ_gt.reshape(num_step, num_agent).transpose(0, 1)
+                    agent_occ_gt[agent_occ_gt == -1] = self.encoder.agent_encoder.grid_size // 2
+                    agent_occ_gt = torch.nn.functional.one_hot(agent_occ_gt, num_classes=self.encoder.agent_encoder.grid_size)
+                    agent_occ = self.attr_tokenizer.pad_square(agent_occ.softmax(-1).detach().cpu().numpy())[0]
+                    agent_occ_gt = self.attr_tokenizer.pad_square(agent_occ_gt.detach().cpu().numpy())[0]
+                    plot_occ_grid(pred['scenario_id'][0],
+                                  agent_occ,
+                                  gt_occ=agent_occ_gt,
+                                  mode='agent',
+                                  save_path=self.save_path,
+                                  prefix=f'training_{self.global_step:06d}_')
+                    pt_occ = pt_occ.reshape(num_step, num_pt, -1).transpose(0, 1)
+                    pt_occ_gt = pt_occ_gt.reshape(num_step, num_pt).transpose(0, 1)
+                    pt_occ_gt[pt_occ_gt == -1] = self.encoder.agent_encoder.grid_size // 2
+                    pt_occ_gt = torch.nn.functional.one_hot(pt_occ_gt, num_classes=self.encoder.agent_encoder.grid_size)
+                    pt_occ = self.attr_tokenizer.pad_square(pt_occ.sigmoid().detach().cpu().numpy())[0]
+                    pt_occ_gt = self.attr_tokenizer.pad_square(pt_occ_gt.detach().cpu().numpy())[0]
+                    plot_occ_grid(pred['scenario_id'][0],
+                                  pt_occ,
+                                  gt_occ=pt_occ_gt,
+                                  mode='pt',
+                                  save_path=self.save_path,
+                                  prefix=f'training_{self.global_step:06d}_')
+
+            return loss
+
+        train_mask = data['agent']['train_mask']
+        # remove_ina_mask = data['agent']['remove_ina_mask']
+
+        # motion token loss
+        if self.predict_motion:
+
+            next_token_idx = pred['next_token_idx']
+            next_token_prob = pred['next_token_prob'] # (a, t, token_size)
+            next_token_idx_gt = pred['next_token_idx_gt'] # (a, t)
+            next_token_eval_mask = pred['next_token_eval_mask'] # (a, t)
+            next_token_eval_mask &= train_mask[:, None]
+
+            token_cls_loss = self.token_cls_loss(next_token_prob[next_token_eval_mask],
+                                                 next_token_idx_gt[next_token_eval_mask]) * self.loss_weight['token_cls_loss']
+            self.log('token_cls_loss', token_cls_loss, **log_params)
+
+            loss = loss + token_cls_loss
+
+            # record motion predict precision of certain timesteps of centain type of agents
+            with torch.no_grad():
+                agent_state_idx_gt = data['agent']['state_idx']
+                index = torch.nonzero(agent_state_idx_gt == self.enter_state)
+                for i in range(10):
+                    index[:, 1] += 1
+                    index = index[index[:, 1] < agent_state_idx_gt.shape[1] - 1]
+                    prob = next_token_prob[index[:, 0], index[:, 1]]
+                    gt = next_token_idx_gt[index[:, 0], index[:, 1]]
+                    mask = next_token_eval_mask[index[:, 0], index[:, 1]]
+                    step_token_cls_loss = self.token_cls_loss(prob[mask], gt[mask])
+                    self.log(f's{i}', step_token_cls_loss, prog_bar=True, on_step=False, on_epoch=True, batch_size=1, sync_dist=True)
+
+        # state token loss
+        if self.predict_state:
+
+            next_state_idx = pred['next_state_idx']
+            next_state_prob = pred['next_state_prob']
+            next_state_idx_gt = pred['next_state_idx_gt']
+            next_state_eval_mask = pred['next_state_eval_mask'] # (num_agent, num_timestep)
+
+            state_cls_loss = self.state_cls_loss(next_state_prob[next_state_eval_mask],
+                                                 next_state_idx_gt[next_state_eval_mask]) * self.loss_weight['state_cls_loss']
+            if torch.isnan(state_cls_loss):
+                print("Found NaN in state_cls_loss!!!")
+                print(next_state_prob.shape)
+                print(next_state_idx_gt.shape)
+                print(next_state_eval_mask.shape)
+                print(next_state_idx_gt[next_state_eval_mask].shape)
+            self.log('state_cls_loss', state_cls_loss, **log_params)
+
+            loss = loss + state_cls_loss
+
+            next_state_idx_seed     = pred['next_state_idx_seed']
+            next_state_prob_seed    = pred['next_state_prob_seed']
+            next_state_idx_gt_seed  = pred['next_state_idx_gt_seed']
+            next_type_prob_seed     = pred['next_type_prob_seed']
+            next_type_idx_gt_seed   = pred['next_type_idx_gt_seed']
+            next_shape_seed         = pred['next_shape_seed']
+            next_shape_gt_seed      = pred['next_shape_gt_seed']
+            next_state_eval_mask_seed = pred['next_state_eval_mask_seed']
+            next_attr_eval_mask_seed = pred['next_attr_eval_mask_seed']
+
+            # when num_seed_gt=0 loss term will be NaN
+            state_cls_loss_seed = self.state_cls_loss_seed(next_state_prob_seed[next_state_eval_mask_seed],
+                                                           next_state_idx_gt_seed[next_state_eval_mask_seed]) * self.loss_weight['state_cls_loss']
+            state_cls_loss_seed = torch.nan_to_num(state_cls_loss_seed)
+            self.log('seed_state_cls_loss', state_cls_loss_seed, **log_params)
+
+            type_cls_loss_seed = self.type_cls_loss_seed(next_type_prob_seed[next_attr_eval_mask_seed],
+                                                         next_type_idx_gt_seed[next_attr_eval_mask_seed]) * self.loss_weight['type_cls_loss']
+            shape_reg_loss_seed = self.shape_reg_loss_seed(next_shape_seed[next_attr_eval_mask_seed],
+                                                           next_shape_gt_seed[next_attr_eval_mask_seed]) * self.loss_weight['shape_reg_loss']
+            type_cls_loss_seed = torch.nan_to_num(type_cls_loss_seed)
+            shape_reg_loss_seed = torch.nan_to_num(shape_reg_loss_seed)
+            self.log('seed_type_cls_loss', type_cls_loss_seed, **log_params)
+            self.log('seed_shape_reg_loss', shape_reg_loss_seed, **log_params)
+
+            loss = loss + state_cls_loss_seed + type_cls_loss_seed + shape_reg_loss_seed
+
+            next_head_rel_prob_seed     = pred['next_head_rel_prob_seed']
+            next_head_rel_index_gt_seed = pred['next_head_rel_index_gt_seed']
+            next_offset_xy_seed         = pred['next_offset_xy_seed']
+            next_offset_xy_gt_seed      = pred['next_offset_xy_gt_seed']
+            next_head_eval_mask_seed    = pred['next_head_eval_mask_seed']
+
+            if self.use_grid_token:
+                next_pos_rel_prob_seed      = pred['next_pos_rel_prob_seed']
+                next_pos_rel_index_gt_seed  = pred['next_pos_rel_index_gt_seed']
+
+                pos_cls_loss_seed = self.pos_cls_loss_seed(next_pos_rel_prob_seed[next_attr_eval_mask_seed],
+                                                           next_pos_rel_index_gt_seed[next_attr_eval_mask_seed]) * self.loss_weight['pos_cls_loss']
+                offset_reg_loss_seed = self.offset_reg_loss_seed(next_offset_xy_seed[next_head_eval_mask_seed],
+                                                                next_offset_xy_gt_seed[next_head_eval_mask_seed]) * self.loss_weight['offset_reg_loss']
+                pos_cls_loss_seed = torch.nan_to_num(pos_cls_loss_seed)
+                self.log('seed_pos_cls_loss', pos_cls_loss_seed, **log_params)
+                self.log('seed_offset_reg_loss', offset_reg_loss_seed, **log_params)
+
+                loss = loss + pos_cls_loss_seed + offset_reg_loss_seed
+
+            else:
+                next_pos_rel_xy_seed       = pred['next_pos_rel_xy_seed']
+                next_pos_rel_xy_gt_seed    = pred['next_pos_rel_xy_gt_seed']
+                pos_reg_loss_seed = self.pos_reg_loss_seed(next_pos_rel_xy_seed[next_attr_eval_mask_seed],
+                                                           next_pos_rel_xy_gt_seed[next_attr_eval_mask_seed]) * self.loss_weight['pos_reg_loss']
+                pos_reg_loss_seed = torch.nan_to_num(pos_reg_loss_seed)
+                self.log('seed_pos_reg_loss', pos_reg_loss_seed, **log_params)
+                loss = loss + pos_reg_loss_seed
+
+            head_cls_loss_seed = self.head_cls_loss_seed(next_head_rel_prob_seed[next_head_eval_mask_seed],
+                                                         next_head_rel_index_gt_seed[next_head_eval_mask_seed]) * self.loss_weight['head_cls_loss']
+            self.log('seed_head_cls_loss', head_cls_loss_seed, **log_params)
+
+            loss = loss + head_cls_loss_seed
+
+            # plot_scenario_ids = ['74ad7b76d5906d39', '1351ea8b8333ddcb', '1352066cc3c0508d', '135436833ce5b9e7', '13570a32432d449', '13577c32a81336fb']
+            if random.random() < 4e-5 or int(os.getenv('DEBUG', 0)):
+                with torch.no_grad():
+                    # plot probability of inserting new agent (agent-timestep)
+                    raw_next_state_prob_seed = pred['raw_next_state_prob_seed']
+                    plot_prob_seed(pred['scenario_id'][0],
+                                   torch.softmax(raw_next_state_prob_seed, dim=-1
+                                                 )[..., -1].detach().cpu().numpy(),
+                                   self.save_path,
+                                   prefix=f'training_{self.global_step:06d}_',
+                                   indices=pred['target_indices'].cpu().numpy())
+
+                    # plot heatmap of inserting new agent
+                    if self.use_grid_token:
+                        next_pos_rel_prob_seed = pred['next_pos_rel_prob_seed']
+                        if next_pos_rel_prob_seed.shape[0] > 0:
+                            next_pos_rel_prob_seed = torch.softmax(next_pos_rel_prob_seed, dim=-1).detach().cpu().numpy()
+                            indices = next_pos_rel_index_gt_seed.detach().cpu().numpy()
+                            mask = next_attr_eval_mask_seed.detach().cpu().numpy().astype(np.bool_)
+                            indices[~mask] = -1
+                            prob, indices = self.attr_tokenizer.pad_square(next_pos_rel_prob_seed, indices)
+                            plot_insert_grid(pred['scenario_id'][0],
+                                            prob,
+                                            indices=indices,
+                                            save_path=self.save_path,
+                                            prefix=f'training_{self.global_step:06d}_')
+
+        if self.predict_occ:
+
+            neighbor_agent_grid_idx             = pred['neighbor_agent_grid_idx']
+            neighbor_agent_grid_index_gt        = pred['neighbor_agent_grid_index_gt']
+            neighbor_agent_grid_index_eval_mask = pred['neighbor_agent_grid_index_eval_mask']
+            neighbor_pt_grid_idx                = pred['neighbor_pt_grid_idx']
+            neighbor_pt_grid_index_gt           = pred['neighbor_pt_grid_index_gt']
+            neighbor_pt_grid_index_eval_mask    = pred['neighbor_pt_grid_index_eval_mask']
+
+            neighbor_agent_grid_cls_loss = self.occ_cls_loss(neighbor_agent_grid_idx[neighbor_agent_grid_index_eval_mask],
+                                                             neighbor_agent_grid_index_gt[neighbor_agent_grid_index_eval_mask])
+            neighbor_pt_grid_cls_loss = self.occ_cls_loss(neighbor_pt_grid_idx[neighbor_pt_grid_index_eval_mask],
+                                                          neighbor_pt_grid_index_gt[neighbor_pt_grid_index_eval_mask])
+            # self.log('neighbor_agent_grid_cls_loss', neighbor_agent_grid_cls_loss, **log_params)
+            # self.log('neighbor_pt_grid_cls_loss', neighbor_pt_grid_cls_loss, **log_params)
+            # loss = loss + neighbor_agent_grid_cls_loss + neighbor_pt_grid_cls_loss
+
+            grid_agent_occ_seed = pred['grid_agent_occ_seed']
+            grid_pt_occ_seed = pred['grid_pt_occ_seed']
+            grid_agent_occ_gt_seed = pred['grid_agent_occ_gt_seed'].float()
+            grid_pt_occ_gt_seed = pred['grid_pt_occ_gt_seed'].float()
+            grid_agent_occ_eval_mask_seed = pred['grid_agent_occ_eval_mask_seed']
+            grid_pt_occ_eval_mask_seed = pred['grid_pt_occ_eval_mask_seed']
+
+            # plot_scenario_ids = ['74ad7b76d5906d39', '1351ea8b8333ddcb', '1352066cc3c0508d', '135436833ce5b9e7', '13570a32432d449', '13577c32a81336fb']
+            if random.random() < 4e-5 or os.getenv('DEBUG'):
+                with torch.no_grad():
+                    agent_occ = self.attr_tokenizer.pad_square(grid_agent_occ_seed.sigmoid().detach().cpu().numpy())[0]
+                    agent_occ_gt = self.attr_tokenizer.pad_square(grid_agent_occ_gt_seed.detach().cpu().numpy())[0]
+                    plot_occ_grid(pred['scenario_id'][0],
+                                 agent_occ,
+                                 gt_occ=agent_occ_gt,
+                                 mode='agent',
+                                 save_path=self.save_path,
+                                 prefix=f'training_{self.global_step:06d}_')
+                    pt_occ = self.attr_tokenizer.pad_square(grid_pt_occ_seed.sigmoid().detach().cpu().numpy())[0]
+                    pt_occ_gt = self.attr_tokenizer.pad_square(grid_pt_occ_gt_seed.detach().cpu().numpy())[0]
+                    plot_occ_grid(pred['scenario_id'][0],
+                                 pt_occ,
+                                 gt_occ=pt_occ_gt,
+                                 mode='pt',
+                                 save_path=self.save_path,
+                                 prefix=f'training_{self.global_step:06d}_')
+
+            grid_agent_occ_gt_seed[grid_agent_occ_gt_seed == -1] = 0
+            if grid_agent_occ_gt_seed.min() < 0 or grid_agent_occ_gt_seed.max() > 1 or \
+               grid_pt_occ_gt_seed.min() < 0 or grid_pt_occ_gt_seed.max() > 1:
+                   raise RuntimeError("Occurred invalid values in occ gt")
+
+            agent_occ_loss = self.agent_occ_loss_seed(grid_agent_occ_seed[grid_agent_occ_eval_mask_seed],
+                                                      grid_agent_occ_gt_seed[grid_agent_occ_eval_mask_seed]) * self.loss_weight['agent_occ_loss']
+            pt_occ_loss = self.pt_occ_loss_seed(grid_pt_occ_seed[grid_pt_occ_eval_mask_seed],
+                                                grid_pt_occ_gt_seed[grid_pt_occ_eval_mask_seed]) * self.loss_weight['pt_occ_loss']
+
+            self.log('agent_occ_loss', agent_occ_loss, **log_params)
+            self.log('pt_occ_loss', pt_occ_loss, **log_params)
+            loss = loss + agent_occ_loss + pt_occ_loss
+
+        if os.getenv('LOG_TRAIN', False) and (self.predict_motion or self.predict_state):
+            for a in range(next_token_idx.shape[0]):
+                print(f"agent: {a}")
+                if self.predict_motion:
+                    print(f"pred motion: {next_token_idx[a, :, 0].tolist()}, \ngt motion:   {next_token_idx_gt[a, :].tolist()}")
+                    print(f"train mask: {next_token_eval_mask[a].long().tolist()}")
+                if self.predict_state:
+                    print(f"pred state: {next_state_idx[a, :, 0].tolist()}, \ngt state:   {next_state_idx_gt[a, :].tolist()}")
+                    print(f"train mask: {next_state_eval_mask[a].long().tolist()}")
+            num_sa = next_state_idx_seed[..., 0].sum(dim=-1).bool().sum()
+            for sa in range(num_sa):
+                print(f"seed agent: {sa}")
+                print(f"seed pred state: {next_state_idx_seed[sa, :, 0].tolist()}, \ngt seed state:   {next_state_idx_gt_seed[sa, :].tolist()}")
+                # if sa < next_pos_rel_seed.shape[0]:
+                #     print(f"pred pos: {next_pos_rel_seed[sa, :, 0].tolist()}, \ngt pos:   {next_pos_rel_gt_seed[sa, :, 0].tolist()}")
+                #     print(f"pred head: {next_head_rel_seed[sa].tolist()}, \ngt head:   {next_head_rel_gt_seed[sa].tolist()}")
+                # print(f"seed train mask: {next_state_eval_mask_seed[sa].long().tolist()}")
+
+        # map token loss
+        if self.predict_map:
+
+            map_next_token_prob = pred['map_next_token_prob']
+            map_next_token_idx_gt = pred['map_next_token_idx_gt']
+            map_next_token_eval_mask = pred['map_next_token_eval_mask']
+
+            map_token_loss = self.map_token_loss(map_next_token_prob[map_next_token_eval_mask], map_next_token_idx_gt[map_next_token_eval_mask]) * self.loss_weight['map_token_loss']
+            self.log('map_token_loss', map_token_loss, prog_bar=True, on_step=True, on_epoch=True, batch_size=1)
+            loss = loss + map_token_loss
+
+        allocated = torch.cuda.memory_allocated(device='cuda:0') / (1024 ** 3)
+        reserved = torch.cuda.memory_reserved(device='cuda:0') / (1024 ** 3)
+        self.log('allocated', allocated, **log_params)
+        self.log('reserved', reserved, **log_params)
+
+        return loss
+
+    def validation_step(self,
+                        data,
+                        batch_idx):
+
+        self.debug = int(os.getenv('DEBUG', 0))
+
+        # ! validation in training process
+        if (
+            self._mode == 'training' and (
+            self.current_epoch not in [5, 10, 20, 25, self.max_epochs] or random.random() > 5e-4) and
+            not self.debug
+        ):
+            self.val_open_loop = False
+            self.val_close_loop = False
+            return
+
+        if int(os.getenv('NO_VAL', 0)) or int(os.getenv("CHECK_INPUTS", 0)):
+            return
+
+        # ! check if save exists
+        if not self._plot_rollouts:
+            rollouts_path = os.path.join(self.save_path, f'idx_{self.trainer.global_rank}_{batch_idx}_rollouts.pkl')
+            if os.path.exists(rollouts_path):
+                tqdm.write(f'Skipped batch {batch_idx}')
+                return
+        else:
+            rollouts_path = os.path.join(self.save_path, f'{data["scenario_id"][0]}.gif')
+            if os.path.exists(rollouts_path):
+                tqdm.write(f'Skipped scenario {data["scenario_id"][0]}')
+                return
+
+        # ! data preparation
+        data = self.token_processer(data)
+
+        data = self.match_token_map(data)
+        data = self.sample_pt_pred(data)
+
+        # find map tokens for entering agents
+        data = self._fetch_enterings(data)
+
+        data['batch_size_a'] = data['agent']['ptr'][1:] - data['agent']['ptr'][:-1]
+        data['batch_size_pl'] = data['pt_token']['ptr'][1:] - data['pt_token']['ptr'][:-1]
+        if isinstance(data, Batch):
+            data['agent']['av_index'] += data['agent']['ptr'][:-1]
+
+        if int(os.getenv('NEAREST_POS', 0)):
+            pred = self.encoder.predict_nearest_pos(data, rank=self.local_rank)
+            return
+
+        # if self.insert_agent:
+        #     pred = self.encoder.insert_agent(data)
+        #     return
+
+        # ! open-loop validation
+        if self.val_open_loop or int(os.getenv('OPEN_LOOP', 0)):
+
+            pred = self(data)
+
+            # pred['next_state_prob_seed'] = torch.softmax(pred['next_state_prob_seed'], dim=-1)[..., -1]
+            # plot_prob_seed(pred, self.save_path, suffix=f'_training')
+
+            loss = 0
+
+            if self.predict_motion:
+
+                # motion token
+                next_token_idx = pred['next_token_idx']
+                next_token_idx_gt = pred['next_token_idx_gt'] # (num_agent, num_step, 10)
+                next_token_prob = pred['next_token_prob']
+                next_token_eval_mask = pred['next_token_eval_mask']
+
+                token_cls_loss = self.token_cls_loss(next_token_prob[next_token_eval_mask], next_token_idx_gt[next_token_eval_mask])
+                loss = loss + token_cls_loss
+
+            if self.predict_state:
+
+                # state token
+                next_state_idx = pred['next_state_idx']
+                next_state_idx_gt = pred['next_state_idx_gt']
+                next_state_prob = pred['next_state_prob']
+                next_state_eval_mask = pred['next_state_eval_mask']
+
+                state_cls_loss = self.state_cls_loss(next_state_prob[next_state_eval_mask], next_state_idx_gt[next_state_eval_mask])
+                loss = loss + state_cls_loss
+
+                # seed state token
+                next_state_idx_seed = pred['next_state_idx_seed']
+                next_state_idx_gt_seed = pred['next_state_idx_gt_seed']
+
+            if self.predict_occ:
+
+                grid_agent_occ_seed = pred['grid_agent_occ_seed']
+                grid_pt_occ_seed = pred['grid_pt_occ_seed']
+                grid_agent_occ_gt_seed = pred['grid_agent_occ_gt_seed'].float()
+                grid_pt_occ_gt_seed = pred['grid_pt_occ_gt_seed'].float()
+
+                agent_occ = self.attr_tokenizer.pad_square(grid_agent_occ_seed.sigmoid().detach().cpu().numpy())[0]
+                agent_occ_gt = self.attr_tokenizer.pad_square(grid_agent_occ_gt_seed.detach().cpu().numpy())[0]
+                plot_occ_grid(pred['scenario_id'][0],
+                              agent_occ,
+                              gt_occ=agent_occ_gt,
+                              mode='agent',
+                              save_path=self.save_path,
+                              prefix=f'eval_')
+                pt_occ = self.attr_tokenizer.pad_square(grid_pt_occ_seed.sigmoid().detach().cpu().numpy())[0]
+                pt_occ_gt = self.attr_tokenizer.pad_square(grid_pt_occ_gt_seed.detach().cpu().numpy())[0]
+                plot_occ_grid(pred['scenario_id'][0],
+                              pt_occ,
+                              gt_occ=pt_occ_gt,
+                              mode='pt',
+                              save_path=self.save_path,
+                              prefix=f'eval_')
+
+            self.log('val_loss', loss, prog_bar=True, on_step=True, on_epoch=True, batch_size=1, sync_dist=True)
+
+        if self.val_insert:
+
+            pred = self(data)
+
+            next_state_idx_seed = pred['next_state_idx_seed']
+            next_state_idx_gt_seed = pred['next_state_idx_gt_seed']
+
+            self.NumInsertAccuracy.update(next_state_idx_seed=next_state_idx_seed,
+                                          next_state_idx_gt_seed=next_state_idx_gt_seed)
+
+            return
+
+        # ! close-loop validation
+        if self.val_close_loop and (self.predict_motion or self.predict_state):
+
+            rollouts = []
+            for _ in tqdm(range(self.n_rollout_close_val), leave=False, desc='Rollout ...'):
+                rollout = self.encoder.inference(data.clone())
+            rollouts.append(rollout)
+
+            av_index = int(rollout['ego_index'])
+            scenario_id = rollout['scenario_id'][0]
+
+            # motion tokens
+            if self.predict_motion:
+
+                if self._plot_rollouts:  # only plot gifs for last 2 epochs for efficiency
+                    plot_val(data, rollout, av_index, self.save_path, pl2seed_radius=self.pl2seed_radius, attr_tokenizer=self.attr_tokenizer)
+
+                # next_token_idx = pred['next_token_idx'][..., None]
+                # next_token_idx_gt = pred['next_token_idx_gt'][:, 2:] # hard code 2=11//5
+                # next_token_eval_mask = pred['next_token_eval_mask'][:, 2:]
+
+                # gt_traj = pred['gt_traj']
+                # pred_traj = pred['pred_traj']
+                # pred_head = pred['pred_head']
+
+                # self.TokenCls.update(pred=next_token_idx[next_token_eval_mask], target=next_token_idx_gt[next_token_eval_mask],
+                #                      valid_mask=next_token_eval_mask[next_token_eval_mask])
+                # self.log('val_token_cls_acc', self.TokenCls, prog_bar=True, on_step=True, on_epoch=True, batch_size=1, sync_dist=True)
+
+                # remove the agents which are unseen at current step
+                # eval_mask = data['agent']['valid_mask'][:, self.num_historical_steps - 1]
+
+                # self.minADE.update(pred=pred_traj[eval_mask], target=gt_traj[eval_mask], valid_mask=valid_mask[eval_mask])
+                # self.minFDE.update(pred=pred_traj[eval_mask], target=gt_traj[eval_mask], valid_mask=valid_mask[eval_mask])
+                # print('ade: ', self.minADE.compute(), 'fde: ', self.minFDE.compute())
+
+                # self.log('val_minADE', self.minADE, prog_bar=True, on_step=True, on_epoch=True, batch_size=1)
+                # self.log('val_minFDE', self.minFDE, prog_bar=True, on_step=True, on_epoch=True, batch_size=1)
+
+            # state tokens
+            if self.predict_state:
+
+                if self.use_grid_token:
+                    next_pos_rel_prob_seed = rollout['next_pos_rel_prob_seed'].cpu().numpy() # (s, t, grid_size)
+                    prob, _ = self.attr_tokenizer.pad_square(next_pos_rel_prob_seed)
+
+                if self._plot_rollouts:
+                    if self.use_grid_token:
+                        plot_insert_grid(scenario_id,
+                                        prob,
+                                        save_path=self.save_path,
+                                        prefix=f'inference_')
+                    plot_prob_seed(scenario_id,
+                                   rollout['next_state_prob_seed'].cpu().numpy(),
+                                   self.save_path,
+                                   prefix=f'inference_')
+
+                next_state_idx = rollout['next_state_idx'][..., None]
+                # next_state_idx_gt = rollout['next_state_idx_gt'][:, 2:]
+                # next_state_eval_mask = rollout['next_state_eval_mask'][:, 2:]
+
+                # self.StateCls.update(pred=next_state_idx[next_token_eval_mask], target=next_state_idx_gt[next_token_eval_mask],
+                #                      valid_mask=next_token_eval_mask[next_token_eval_mask])
+                # self.log('val_state_cls_acc', self.TokenCls, prog_bar=True, on_step=True, on_epoch=True, batch_size=1, sync_dist=True)
+
+                self.StateAccuracy.update(state_idx=next_state_idx[..., 0])
+                self.log('valid_accuracy', self.StateAccuracy.compute()['valid'], prog_bar=True, on_step=True, on_epoch=True, batch_size=1)
+                self.log('invalid_accuracy', self.StateAccuracy.compute()['invalid'], prog_bar=True, on_step=True, on_epoch=True, batch_size=1)
+                self.local_logger.info(rollout['log_message'])
+                # print(rollout['log_message'])
+                # print(self.StateAccuracy)
+
+            if self.predict_occ:
+
+                grid_agent_occ_seed = rollout['grid_agent_occ_seed']
+                grid_pt_occ_seed = rollout['grid_pt_occ_seed']
+                grid_agent_occ_gt_seed = rollout['grid_agent_occ_gt_seed']
+
+                agent_occ = self.attr_tokenizer.pad_square(grid_agent_occ_seed.sigmoid().cpu().numpy())[0]
+                agent_occ_gt = self.attr_tokenizer.pad_square(grid_agent_occ_gt_seed.sigmoid().cpu().numpy())[0]
+                if self._plot_rollouts:
+                    plot_occ_grid(scenario_id,
+                                  agent_occ,
+                                  gt_occ=agent_occ_gt,
+                                  mode='agent',
+                                  save_path=self.save_path,
+                                  prefix=f'inference_')
+
+            if self._online_metric or self._save_validate_reuslts:
+
+                # ! format results
+                pred_valid, token_pos, token_head = [], [], []
+                pred_traj, pred_head, pred_z = [], [], []
+                pred_shape, pred_type, pred_state = [], [], []
+                agent_id = []
+                for rollout in rollouts:
+                    pred_valid.append(rollout['pred_valid'])
+                    token_pos.append(rollout['pos_a'])
+                    token_head.append(rollout['head_a'])
+                    pred_traj.append(rollout['pred_traj'])
+                    pred_head.append(rollout['pred_head'])
+                    pred_z.append(rollout['pred_z'])
+                    pred_shape.append(rollout['eval_shape'])
+                    pred_type.append(rollout['pred_type'])
+                    pred_state.append(rollout['next_state_idx'])
+                    agent_id.append(rollout['agent_id'])
+
+                pred_valid = torch.stack(pred_valid, dim=1)
+                token_pos = torch.stack(token_pos, dim=1)
+                token_head = torch.stack(token_head, dim=1)
+                pred_traj = torch.stack(pred_traj, dim=1)  # (n_agent, n_rollout, n_step, 2)
+                pred_head = torch.stack(pred_head, dim=1)
+                pred_z = torch.stack(pred_z, dim=1)
+                pred_shape = torch.stack(pred_shape, dim=1)  # [n_agent, n_rollout, 3]
+                pred_type = torch.stack(pred_type, dim=1)  # [n_agent, n_rollout]
+                pred_state = torch.stack(pred_state, dim=1)  # [n_agent, n_rollout, n_step // shift]
+                agent_id = torch.stack(agent_id, dim=1)  # [n_agent, n_rollout]
+
+                agent_batch = torch.zeros((pred_traj.shape[0]), dtype=torch.long)
+                rollouts = dict(
+                    _scenario_id=data['scenario_id'],
+                    scenario_id=get_scenario_id_int_tensor(data['scenario_id']),
+                    av_id=int(rollouts[0]['agent_id'][rollouts[0]['ego_index']]),  # NOTE: hard code!!!
+                    agent_id=agent_id.cpu(),
+                    agent_batch=agent_batch.cpu(),
+                    pred_traj=pred_traj.cpu(),
+                    pred_z=pred_z.cpu(),
+                    pred_head=pred_head.cpu(),
+                    pred_shape=pred_shape.cpu(),
+                    pred_type=pred_type.cpu(),
+                    pred_state=pred_state.cpu(),
+                    pred_valid=pred_valid.cpu(),
+                    token_pos=token_pos.cpu(),
+                    token_head=token_head.cpu(),
+                    tfrecord_path=data['tfrecord_path'],
+                )
+
+                if self._save_validate_reuslts:
+                    with open(rollouts_path, 'wb') as f:
+                        pickle.dump(rollouts, f)
+
+                if self._online_metric:
+                    self._long_metrics.update(rollouts)
+
+    def on_validation_start(self):
+        self.scenario_rollouts = []
+        self.batch_metric = defaultdict(list)
+
+    def on_validation_epoch_end(self):
+        if self.val_close_loop:
+
+            if self._long_metrics is not None:
+                epoch_long_metrics = self._long_metrics.compute()
+                if self.global_rank == 0:
+                    epoch_long_metrics['epoch'] = self.current_epoch
+                    self.logger.log_metrics(epoch_long_metrics)
+
+                self._long_metrics.reset()
+
+            self.minADE.reset()
+            self.minFDE.reset()
+            self.StateAccuracy.reset()
+
+    def configure_optimizers(self):
+        optimizer = torch.optim.AdamW(self.parameters(), lr=self.lr)
+
+        def lr_lambda(current_step):
+            if current_step + 1 < self.warmup_steps:
+                return float(current_step + 1) / float(max(1, self.warmup_steps))
+            return max(
+                0.0, 0.5 * (1.0 + math.cos(math.pi * (current_step - self.warmup_steps) / float(max(1, self.total_steps - self.warmup_steps))))
+            )
+
+        lr_scheduler = LambdaLR(optimizer, lr_lambda=lr_lambda)
+        return [optimizer], [lr_scheduler]
+
+    def load_state_from_file(self, filename, to_cpu=False):
+        logger = self.local_logger
+
+        if not os.path.isfile(filename):
+            raise FileNotFoundError
+
+        logger.info('==> Loading parameters from checkpoint %s to %s' % (filename, 'CPU' if to_cpu else 'GPU'))
+        loc_type = torch.device('cpu') if to_cpu else None
+        checkpoint = torch.load(filename, map_location=loc_type)
+
+        version = checkpoint.get("version", None)
+        if version is not None:
+            logger.info('==> Checkpoint trained from version: %s' % version)
+
+
+        model_state_disk = checkpoint['state_dict']
+        logger.info(f'The number of disk ckpt keys: {len(model_state_disk)}')
+
+        model_state = self.state_dict()
+        model_state_disk_filter = {}
+        for key, val in model_state_disk.items():
+            if key in model_state and model_state_disk[key].shape == model_state[key].shape:
+                model_state_disk_filter[key] = val
+            else:
+                if key not in model_state:
+                    print(f'Ignore key in disk (not found in model): {key}, shape={val.shape}')
+                else:
+                    print(f'Ignore key in disk (shape does not match): {key}, load_shape={val.shape}, model_shape={model_state[key].shape}')
+
+        model_state_disk = model_state_disk_filter
+        missing_keys, unexpected_keys = self.load_state_dict(model_state_disk, strict=False)
+
+        logger.info(f'Missing keys: {missing_keys}')
+        logger.info(f'The number of missing keys: {len(missing_keys)}')
+        logger.info(f'The number of unexpected keys: {len(unexpected_keys)}')
+        logger.info('==> Done (total keys %d)' % (len(model_state)))
+
+        epoch = checkpoint.get('epoch', -1)
+        it = checkpoint.get('it', 0.0)
+
+        return it, epoch
+
+    def match_token_map(self, data):
+        traj_pos = data['map_save']['traj_pos'].to(torch.float)
+        traj_theta = data['map_save']['traj_theta'].to(torch.float)
+        pl_idx_list = data['map_save']['pl_idx_list']
+        token_sample_pt = self.map_token['sample_pt'].to(traj_pos.device)
+        token_src = self.map_token['traj_src'].to(traj_pos.device)
+        max_traj_len = self.map_token['traj_src'].shape[1]
+        pl_num = traj_pos.shape[0]
+
+        pt_token_pos = traj_pos[:, 0, :].clone()
+        pt_token_orientation = traj_theta.clone()
+        cos, sin = traj_theta.cos(), traj_theta.sin()
+        rot_mat = traj_theta.new_zeros(pl_num, 2, 2)
+        rot_mat[..., 0, 0] = cos
+        rot_mat[..., 0, 1] = -sin
+        rot_mat[..., 1, 0] = sin
+        rot_mat[..., 1, 1] = cos
+        traj_pos_local = torch.bmm((traj_pos - traj_pos[:, 0:1]), rot_mat.view(-1, 2, 2))
+        distance = torch.sum((token_sample_pt[None] - traj_pos_local.unsqueeze(1)) ** 2, dim=(-2, -1))
+        pt_token_id = torch.argmin(distance, dim=1)
+
+        if self.noise:
+            topk_indices = torch.argsort(torch.sum((token_sample_pt[None] - traj_pos_local.unsqueeze(1))**2, dim=(-2, -1)), dim=1)[:, :8]
+            sample_topk = torch.randint(0, topk_indices.shape[-1], size=(topk_indices.shape[0], 1), device=topk_indices.device)
+            pt_token_id = torch.gather(topk_indices, 1, sample_topk).squeeze(-1)
+
+        # cos, sin = traj_theta.cos(), traj_theta.sin()
+        # rot_mat = traj_theta.new_zeros(pl_num, 2, 2)
+        # rot_mat[..., 0, 0] = cos
+        # rot_mat[..., 0, 1] = sin
+        # rot_mat[..., 1, 0] = -sin
+        # rot_mat[..., 1, 1] = cos
+        # token_src_world = torch.bmm(token_src[None, ...].repeat(pl_num, 1, 1, 1).reshape(pl_num, -1, 2),
+        #                             rot_mat.view(-1, 2, 2)).reshape(pl_num, token_src.shape[0], max_traj_len, 2) + traj_pos[:, None, [0], :]
+        # token_src_world_select = token_src_world.view(-1, 1024, 11, 2)[torch.arange(pt_token_id.view(-1).shape[0]), pt_token_id.view(-1)].view(pl_num, max_traj_len, 2)
+
+        pl_idx_full = pl_idx_list.clone()
+        token2pl = torch.stack([torch.arange(len(pl_idx_list), device=traj_pos.device), pl_idx_full.long()])
+        count_nums = []
+        for pl in pl_idx_full.unique():
+            pt = token2pl[0, token2pl[1, :] == pl]
+            left_side = (data['pt_token']['side'][pt] == 0).sum()
+            right_side = (data['pt_token']['side'][pt] == 1).sum()
+            center_side = (data['pt_token']['side'][pt] == 2).sum()
+            count_nums.append(torch.Tensor([left_side, right_side, center_side]))
+        count_nums = torch.stack(count_nums, dim=0)
+        num_polyline = int(count_nums.max().item())
+        traj_mask = torch.zeros((int(len(pl_idx_full.unique())), 3, num_polyline), dtype=bool)
+        idx_matrix = torch.arange(traj_mask.size(2)).unsqueeze(0).unsqueeze(0)
+        idx_matrix = idx_matrix.expand(traj_mask.size(0), traj_mask.size(1), -1)
+        counts_num_expanded = count_nums.unsqueeze(-1)
+        mask_update = idx_matrix < counts_num_expanded
+        traj_mask[mask_update] = True
+
+        data['pt_token']['traj_mask'] = traj_mask
+        data['pt_token']['position'] = torch.cat([pt_token_pos, torch.zeros((data['pt_token']['num_nodes'], 1),
+                                                                            device=traj_pos.device, dtype=torch.float)], dim=-1)
+        data['pt_token']['orientation'] = pt_token_orientation
+        data['pt_token']['height'] = data['pt_token']['position'][:, -1]
+        data[('pt_token', 'to', 'map_polygon')] = {}
+        data[('pt_token', 'to', 'map_polygon')]['edge_index'] = token2pl  # (2, num_points)
+        data['pt_token']['token_idx'] = pt_token_id
+
+        # data['pt_token']['batch'] = torch.zeros(data['pt_token']['num_nodes'], device=traj_pos.device).long()
+        # data['pt_token']['ptr'] = torch.tensor([0, data['pt_token']['num_nodes']], device=traj_pos.device).long()
+
+        return data
+
+    def sample_pt_pred(self, data):
+        traj_mask = data['pt_token']['traj_mask']
+        raw_pt_index = torch.arange(1, traj_mask.shape[2]).repeat(traj_mask.shape[0], traj_mask.shape[1], 1)
+        masked_pt_index = raw_pt_index.view(-1)[torch.randperm(raw_pt_index.numel())[:traj_mask.shape[0] * traj_mask.shape[1] * ((traj_mask.shape[2] - 1) // 3)].reshape(traj_mask.shape[0], traj_mask.shape[1], (traj_mask.shape[2] - 1) // 3)]
+        masked_pt_index = torch.sort(masked_pt_index, -1)[0]
+        pt_valid_mask = traj_mask.clone()
+        pt_valid_mask.scatter_(2, masked_pt_index, False)
+        pt_pred_mask = traj_mask.clone()
+        pt_pred_mask.scatter_(2, masked_pt_index, False)
+        tmp_mask = pt_pred_mask.clone()
+        tmp_mask[:, :, :] = True
+        tmp_mask.scatter_(2, masked_pt_index-1, False)
+        pt_pred_mask.masked_fill_(tmp_mask, False)
+        pt_pred_mask = pt_pred_mask * torch.roll(traj_mask, shifts=-1, dims=2)
+        pt_target_mask = torch.roll(pt_pred_mask, shifts=1, dims=2)
+
+        data['pt_token']['pt_valid_mask'] = pt_valid_mask[traj_mask]
+        data['pt_token']['pt_pred_mask'] = pt_pred_mask[traj_mask]
+        data['pt_token']['pt_target_mask'] = pt_target_mask[traj_mask]
+
+        return data
+
+    def _fetch_enterings(self, data: HeteroData, plot: bool=False):
+        data['agent']['grid_token_idx']     = torch.zeros_like(data['agent']['state_idx']).long()
+        data['agent']['grid_offset_xy']     = torch.zeros_like(data['agent']['token_pos'])
+        data['agent']['heading_token_idx']  = torch.zeros_like(data['agent']['state_idx']).long()
+        data['agent']['sort_indices']       = torch.zeros_like(data['agent']['state_idx']).long()
+        data['agent']['inrange_mask']       = torch.zeros_like(data['agent']['state_idx']).bool()
+        data['agent']['bos_mask']           = torch.zeros_like(data['agent']['state_idx']).bool()
+
+        data['agent']['pos_xy']     = torch.zeros_like(data['agent']['token_pos'])
+        if self.predict_occ:
+            num_step = data['agent']['state_idx'].shape[1]
+            data['agent']['pt_grid_token_idx'] = torch.zeros_like(data['pt_token']['token_idx'])[None].repeat(num_step, 1).long()
+
+        for b in range(data.num_graphs):
+            av_index = int(data['agent']['av_index'][b])
+            agent_batch_mask = data['agent']['batch'] == b
+            pt_batch_mask = data['pt_token']['batch'] == b
+            pt_token_idx = data['pt_token']['token_idx'][pt_batch_mask]
+            pt_pos = data['pt_token']['position'][pt_batch_mask]
+            agent_token_pos = data['agent']['token_pos'][agent_batch_mask]
+            agent_token_heading = data['agent']['token_heading'][agent_batch_mask]
+            state_idx = data['agent']['state_idx'][agent_batch_mask]
+            ego_pos = agent_token_pos[av_index]  # NOTE: `av_index` will be added by `ptr` later
+            ego_heading = agent_token_heading[av_index]
+
+            grid_token_idx = torch.full(state_idx.shape, -1, device=state_idx.device)
+            offset_xy = torch.zeros_like(agent_token_pos)
+            sort_indices = torch.zeros_like(grid_token_idx)
+            pt_grid_token_idx = torch.full((state_idx.shape[1], *pt_token_idx.shape), -1, device=pt_token_idx.device)
+
+            pos_xy = torch.zeros((*state_idx.shape, 2), device=state_idx.device)
+
+            is_bos = []
+            is_inrange = []
+            for t in range(agent_token_pos.shape[1]): # num_step
+
+                # tokenize position
+                is_bos_t = state_idx[:, t] == self.enter_state
+                is_invalid_t = state_idx[:, t] == self.invalid_state
+                is_inrange_t = ((agent_token_pos[:, t] - ego_pos[[t]]) ** 2).sum(-1).sqrt() <= self.pl2seed_radius
+                grid_index_t, offset_xy_t = self.attr_tokenizer.encode_pos(x=agent_token_pos[~is_invalid_t & is_inrange_t, t],
+                                                                           y=ego_pos[[t]],
+                                                                           theta_y=ego_heading[[t]])
+                grid_token_idx[~is_invalid_t & is_inrange_t, t] = grid_index_t
+                offset_xy[~is_invalid_t & is_inrange_t, t] = offset_xy_t
+
+                pos_xy[~is_invalid_t & is_inrange_t, t] = agent_token_pos[~is_invalid_t & is_inrange_t, t] - ego_pos[[t]]
+
+                # distance = ((agent_token_pos[:, t] - ego_pos[[t]]) ** 2).sum(-1).sqrt()
+                head_vector = torch.stack([ego_heading[[t]].cos(), ego_heading[[t]].sin()], dim=-1)
+                distance = angle_between_2d_vectors(ctr_vector=head_vector,
+                                                    nbr_vector=agent_token_pos[:, t] - ego_pos[[t]])
+                # distance = torch.rand(agent_token_pos.shape[0], device=agent_token_pos.device)
+                distance[~(is_bos_t & is_inrange_t)] = torch.inf
+                sort_dist, sort_indice = distance.sort()
+                sort_indice[torch.isinf(sort_dist)] = av_index
+                sort_indices[:, t] = sort_indice
+
+                is_bos.append(is_bos_t)
+                is_inrange.append(is_inrange_t)
+
+                # tokenize pt token
+                if self.predict_occ:
+                    is_inrange_t = ((pt_pos[:, :2] - ego_pos[None, t]) ** 2).sum(-1).sqrt() <= self.pl2seed_radius
+                    grid_index_t, _ = self.attr_tokenizer.encode_pos(x=pt_pos[is_inrange_t, :2],
+                                                                     y=ego_pos[[t]],
+                                                                     theta_y=ego_heading[[t]])
+
+                    pt_grid_token_idx[t, is_inrange_t] = grid_index_t
+
+            # tokenize heading
+            rel_heading = agent_token_heading - ego_heading[None, ...]
+            heading_token_idx = self.attr_tokenizer.encode_heading(rel_heading)
+
+            data['agent']['grid_token_idx'][agent_batch_mask] = grid_token_idx
+            data['agent']['grid_offset_xy'][agent_batch_mask] = offset_xy
+            data['agent']['pos_xy'][agent_batch_mask] = pos_xy
+            data['agent']['heading_token_idx'][agent_batch_mask] = heading_token_idx
+            data['agent']['sort_indices'][agent_batch_mask] = sort_indices
+            data['agent']['inrange_mask'][agent_batch_mask] = torch.stack(is_inrange, dim=1)
+            data['agent']['bos_mask'][agent_batch_mask] = torch.stack(is_bos, dim=1)
+            if self.predict_occ:
+                data['agent']['pt_grid_token_idx'][:, pt_batch_mask] = pt_grid_token_idx
+
+            plot = False
+            if plot:
+                scenario_id = data['scenario_id'][b]
+                dummy_prob = np.zeros((ego_pos.shape[0], self.attr_tokenizer.grid.shape[0])) + .5
+                indices = grid_token_idx[:, 1:][state_idx[:, 1:] == self.enter_state].reshape(-1).cpu().numpy()
+                dummy_prob, indices = self.attr_tokenizer.pad_square(dummy_prob, indices)
+                # plot_insert_grid(scenario_id, dummy_prob,
+                #                  self.attr_tokenizer.grid.cpu().numpy(),
+                #                  ego_pos.cpu().numpy(),
+                #                  None,
+                #                  save_path=os.path.join(self.save_path, 'vis'),
+                #                  indices=indices[np.newaxis, ...],
+                #                  inference=True,
+                #                  all_t_in_one=True)
+
+                enter_index = [grid_token_idx[:, i][state_idx[:, i] == self.enter_state].tolist()
+                               for i in range(agent_token_pos.shape[1])]
+                agent_labels = [[f'A{i}'] * agent_token_pos.shape[1] for i in range(agent_token_pos.shape[0])]
+                plot_scenario(scenario_id,
+                              data['map_point']['position'].cpu().numpy(),
+                              agent_token_pos.cpu().numpy(),
+                              agent_token_heading.cpu().numpy(),
+                              state_idx.cpu().numpy(),
+                              types=list(map(lambda i: self.encoder.agent_encoder.agent_type[i],
+                                             data['agent']['type'].tolist())),
+                              av_index=av_index,
+                              pl2seed_radius=self.pl2seed_radius,
+                              attr_tokenizer=self.attr_tokenizer,
+                              enter_index=enter_index,
+                              save_gif=False,
+                              save_path=os.path.join(self.save_path, 'vis'),
+                              agent_labels=agent_labels,
+                              tokenized=True)
+
+        return data

seq_10_150_3_3_encode_occ_separate_offsets_bs8_128_no_seqindex_long_lastvalid_ablation_grid/backups/dev/modules/attr_tokenizer.py

@@ -0,0 +1,109 @@
+import torch
+import torch.nn as nn
+import numpy as np
+from dev.utils.func import wrap_angle, angle_between_2d_vectors
+
+
+class Attr_Tokenizer(nn.Module):
+
+    def __init__(self, grid_range, grid_interval, radius, angle_interval):
+        super().__init__()
+        self.grid_range = grid_range
+        self.grid_interval = grid_interval
+        self.radius = radius
+        self.angle_interval = angle_interval
+        self.heading = torch.pi / 2
+        self._prepare_grid()
+
+        self.grid_size = self.grid.shape[0]
+        self.angle_size = int(360. / self.angle_interval)
+
+        assert torch.all(self.grid[self.grid_size // 2] == 0.)
+
+    def _prepare_grid(self):
+        num_grid = int(self.grid_range / self.grid_interval) + 1 # Do not use '//'
+
+        x = torch.linspace(0, num_grid - 1, steps=num_grid)
+        y = torch.linspace(0, num_grid - 1, steps=num_grid)
+        grid_x, grid_y = torch.meshgrid(x, y, indexing='xy')
+        grid = torch.stack([grid_x.flatten(), grid_y.flatten()], dim=-1)  # (n^2, 2)
+        grid = grid.reshape(num_grid, num_grid, 2).flip(dims=[0]).reshape(-1, 2)
+        grid = (grid - x.shape[0] // 2) * self.grid_interval
+
+        distance = (grid ** 2).sum(-1).sqrt()
+        square_mask = ((distance <= self.radius) & (distance >= 0.)) | (distance == 0.)
+        self.register_buffer('grid', grid[square_mask])
+        self.register_buffer('dist', torch.norm(self.grid, p=2, dim=-1))
+        head_vector = torch.stack([torch.tensor(self.heading).cos(), torch.tensor(self.heading).sin()])
+        self.register_buffer('dir', angle_between_2d_vectors(ctr_vector=head_vector.unsqueeze(0),
+                                            nbr_vector=self.grid))  # (-pi, pi]
+
+        self.num_grid = num_grid
+        self.square_mask = square_mask.numpy()
+
+    def _apply_rot(self, x, theta):
+        # x: (b, l, 2) e.g. (num_step, num_agent, 2)
+        # theta: (b,) e.g. (num_step,)
+        cos, sin = theta.cos(), theta.sin()
+        rot_mat = torch.zeros((theta.shape[0], 2, 2), device=theta.device)
+        rot_mat[:, 0, 0] = cos
+        rot_mat[:, 0, 1] = sin
+        rot_mat[:, 1, 0] = -sin
+        rot_mat[:, 1, 1] = cos
+        x = torch.bmm(x, rot_mat)
+        return x
+
+    def pad_square(self, prob, indices=None):
+        # square_mask: bool array of shape (n^2,)
+        # prob: float array of shape (num_step, m)
+        pad_prob = np.zeros((*prob.shape[:-1], self.square_mask.shape[0]))
+        pad_prob[..., self.square_mask] = prob
+
+        square_indices = np.arange(self.square_mask.shape[0])
+        circle_indices = np.concatenate([square_indices[self.square_mask], [-1]])
+        if indices is not None:
+            indices = circle_indices[indices]
+
+        return pad_prob, indices
+
+    def get_grid(self, x, theta=None):
+        x = x.reshape(-1, 2)
+        grid = self.grid[None, ...].to(x.device)
+        if theta is not None:
+            grid = self._apply_rot(grid, (theta - self.heading).expand(x.shape[0]))
+        return x[:, None] + grid
+
+    def encode_pos(self, x, y, theta_y=None):
+        assert x.dim() == y.dim() and x.shape[-1] == 2 and y.shape[-1] == 2, \
+                    f"Invalid input shape x: {x.shape}, y: {y.shape}."
+        centered_x = x - y
+        if theta_y is not None:
+            centered_x = self._apply_rot(centered_x[:, None], -(theta_y - self.heading).expand(x.shape[0]))[:, 0]
+        distance = ((centered_x[:, None] - self.grid.to(x.device)[None, ...]) ** 2).sum(-1).sqrt()
+        index = torch.argmin(distance, dim=-1)
+
+        grid_xy = self.grid[index]
+        offset_xy = centered_x - grid_xy
+
+        return index.long(), offset_xy
+
+    def decode_pos(self, index, y=None, theta_y=None):
+        assert torch.all((index >= 0) & (index < self.grid_size))
+        centered_x = self.grid.to(index.device)[index.long()]
+        if y is not None:
+            if theta_y is not None:
+                centered_x = self._apply_rot(centered_x[:, None], (theta_y - self.heading).expand(centered_x.shape[0]))[:, 0]
+            x = centered_x + y
+            return x.float()
+        return centered_x.float()
+
+    def encode_heading(self, heading):
+        heading = (wrap_angle(heading) + torch.pi) / (2 * torch.pi) * 360
+        index = heading // self.angle_interval
+        return index.long()
+
+    def decode_heading(self, index):
+        assert torch.all(index >= 0) and torch.all(index < (360 / self.angle_interval))
+        angles = index * self.angle_interval - 180
+        angles = angles / 360 * (2 * torch.pi)
+        return angles.float()

seq_10_150_3_3_encode_occ_separate_offsets_bs8_128_no_seqindex_long_lastvalid_ablation_grid/backups/dev/modules/debug.py

@@ -0,0 +1,1439 @@
+from typing import Dict, Mapping, Optional
+import math
+import torch
+import torch.nn as nn
+from torch_cluster import radius, radius_graph
+from torch_geometric.data import HeteroData, Batch
+from torch_geometric.utils import dense_to_sparse, subgraph
+
+from dev.modules.layers import *
+from dev.modules.map_decoder import discretize_neighboring
+from dev.utils.geometry import angle_between_2d_vectors, wrap_angle
+from dev.utils.weight_init import weight_init
+
+
+def cal_polygon_contour(x, y, theta, width, length):
+    left_front_x = x + 0.5 * length * math.cos(theta) - 0.5 * width * math.sin(theta)
+    left_front_y = y + 0.5 * length * math.sin(theta) + 0.5 * width * math.cos(theta)
+    left_front = (left_front_x, left_front_y)
+
+    right_front_x = x + 0.5 * length * math.cos(theta) + 0.5 * width * math.sin(theta)
+    right_front_y = y + 0.5 * length * math.sin(theta) - 0.5 * width * math.cos(theta)
+    right_front = (right_front_x, right_front_y)
+
+    right_back_x = x - 0.5 * length * math.cos(theta) + 0.5 * width * math.sin(theta)
+    right_back_y = y - 0.5 * length * math.sin(theta) - 0.5 * width * math.cos(theta)
+    right_back = (right_back_x, right_back_y)
+
+    left_back_x = x - 0.5 * length * math.cos(theta) - 0.5 * width * math.sin(theta)
+    left_back_y = y - 0.5 * length * math.sin(theta) + 0.5 * width * math.cos(theta)
+    left_back = (left_back_x, left_back_y)
+    polygon_contour = [left_front, right_front, right_back, left_back]
+
+    return polygon_contour
+
+
+class SMARTAgentDecoder(nn.Module):
+
+    def __init__(self,
+                 dataset: str,
+                 input_dim: int,
+                 hidden_dim: int,
+                 num_historical_steps: int,
+                 num_interaction_steps: int,
+                 time_span: Optional[int],
+                 pl2a_radius: float,
+                 pl2seed_radius: float,
+                 a2a_radius: float,
+                 num_freq_bands: int,
+                 num_layers: int,
+                 num_heads: int,
+                 head_dim: int,
+                 dropout: float,
+                 token_data: Dict,
+                 token_size: int,
+                 special_token_index: list=[],
+                 predict_motion: bool=False,
+                 predict_state: bool=False,
+                 predict_map: bool=False,
+                 state_token: Dict[str, int]=None,
+                 seed_size: int=5) -> None:
+
+        super(SMARTAgentDecoder, self).__init__()
+        self.dataset = dataset
+        self.input_dim = input_dim
+        self.hidden_dim = hidden_dim
+        self.num_historical_steps = num_historical_steps
+        self.num_interaction_steps = num_interaction_steps
+        self.time_span = time_span if time_span is not None else num_historical_steps
+        self.pl2a_radius = pl2a_radius
+        self.pl2seed_radius = pl2seed_radius
+        self.a2a_radius = a2a_radius
+        self.num_freq_bands = num_freq_bands
+        self.num_layers = num_layers
+        self.num_heads = num_heads
+        self.head_dim = head_dim
+        self.dropout = dropout
+        self.special_token_index = special_token_index
+        self.predict_motion = predict_motion
+        self.predict_state = predict_state
+        self.predict_map = predict_map
+
+        # state tokens
+        self.state_type = list(state_token.keys())
+        self.state_token = state_token
+        self.invalid_state = int(state_token['invalid'])
+        self.valid_state = int(state_token['valid'])
+        self.enter_state = int(state_token['enter'])
+        self.exit_state = int(state_token['exit'])
+
+        self.seed_state_type = ['invalid', 'enter']
+        self.valid_state_type = ['invalid', 'valid', 'exit']
+
+        input_dim_x_a = 2
+        input_dim_r_t = 4
+        input_dim_r_pt2a = 3
+        input_dim_r_a2a = 3
+        input_dim_token = 8 # tokens: (token_size, 4, 2)
+
+        self.seed_size = seed_size
+
+        self.all_agent_type = ['veh', 'ped', 'cyc', 'background', 'invalid', 'seed']
+        self.seed_agent_type = ['veh', 'ped', 'cyc', 'seed']
+        self.type_a_emb = nn.Embedding(len(self.all_agent_type), hidden_dim)
+        self.shape_emb = MLPLayer(3, hidden_dim, hidden_dim)
+        if self.predict_state:
+            self.state_a_emb = nn.Embedding(len(self.state_type), hidden_dim)
+        self.invalid_shape_value = .1
+        self.motion_gap = 1.
+        self.heading_gap = 1.
+
+        self.x_a_emb = FourierEmbedding(input_dim=input_dim_x_a, hidden_dim=hidden_dim, num_freq_bands=num_freq_bands)
+        self.r_t_emb = FourierEmbedding(input_dim=input_dim_r_t, hidden_dim=hidden_dim, num_freq_bands=num_freq_bands)
+        self.r_pt2a_emb = FourierEmbedding(input_dim=input_dim_r_pt2a, hidden_dim=hidden_dim,
+                                           num_freq_bands=num_freq_bands)
+        self.r_a2a_emb = FourierEmbedding(input_dim=input_dim_r_a2a, hidden_dim=hidden_dim,
+                                          num_freq_bands=num_freq_bands)
+        self.token_emb_veh = MLPEmbedding(input_dim=input_dim_token, hidden_dim=hidden_dim)
+        self.token_emb_ped = MLPEmbedding(input_dim=input_dim_token, hidden_dim=hidden_dim)
+        self.token_emb_cyc = MLPEmbedding(input_dim=input_dim_token, hidden_dim=hidden_dim)
+        self.no_token_emb = nn.Embedding(1, hidden_dim)
+        self.bos_token_emb = nn.Embedding(1, hidden_dim)
+        # FIXME: do we need this???
+        self.token_emb_offset = MLPEmbedding(input_dim=2, hidden_dim=hidden_dim)
+
+        num_inputs = 2
+        if self.predict_state:
+            num_inputs = 3
+        self.fusion_emb = MLPEmbedding(input_dim=self.hidden_dim * num_inputs, hidden_dim=self.hidden_dim)
+
+        self.t_attn_layers = nn.ModuleList(
+            [AttentionLayer(hidden_dim=hidden_dim, num_heads=num_heads, head_dim=head_dim, dropout=dropout,
+                            bipartite=False, has_pos_emb=True) for _ in range(num_layers)]
+        )
+        self.pt2a_attn_layers = nn.ModuleList(
+            [AttentionLayer(hidden_dim=hidden_dim, num_heads=num_heads, head_dim=head_dim, dropout=dropout,
+                            bipartite=True, has_pos_emb=True) for _ in range(num_layers)]
+        )
+        self.a2a_attn_layers = nn.ModuleList(
+            [AttentionLayer(hidden_dim=hidden_dim, num_heads=num_heads, head_dim=head_dim, dropout=dropout,
+                            bipartite=False, has_pos_emb=True) for _ in range(num_layers)]
+        )
+        self.token_size = token_size # 2048
+        # agent motion prediction head
+        self.token_predict_head = MLPLayer(input_dim=hidden_dim, hidden_dim=hidden_dim,
+                                           output_dim=self.token_size)
+        # agent state prediction head
+        if self.predict_state:
+
+            self.seed_feature = nn.Embedding(self.seed_size, self.hidden_dim)
+
+            self.state_predict_head = MLPLayer(input_dim=hidden_dim, hidden_dim=hidden_dim,
+                                               output_dim=len(self.valid_state_type))
+
+            self.seed_head = MLPLayer(input_dim=hidden_dim, hidden_dim=hidden_dim,
+                                      output_dim=hidden_dim)
+
+            self.seed_state_predict_head = MLPLayer(input_dim=hidden_dim, hidden_dim=hidden_dim,
+                                                    output_dim=len(self.seed_state_type))
+            self.seed_type_predict_head = MLPLayer(input_dim=hidden_dim, hidden_dim=hidden_dim,
+                                                   output_dim=len(self.seed_agent_type))
+        # entering token prediction
+        # FIXME: this is just under test!!!
+        # self.bos_pl_predict_head = MLPLayer(input_dim=hidden_dim, hidden_dim=hidden_dim,
+        #                                     output_dim=200)
+        # self.bos_offset_predict_head = MLPLayer(input_dim=hidden_dim, hidden_dim=hidden_dim,
+        #                                         output_dim=2601)
+        self.trajectory_token = token_data['token'] # dict('veh', 'ped', 'cyc') (2048, 4, 2)
+        self.trajectory_token_traj = token_data['traj'] # (2048, 6, 3)
+        self.trajectory_token_all = token_data['token_all'] # (2048, 6, 4, 2)
+        self.apply(weight_init)
+
+        self.shift = 5
+        self.beam_size = 5
+        self.hist_mask = True
+        self.temporal_attn_to_invalid = True
+        self.temporal_attn_seed = False
+
+        # FIXME: This is just under test!!!
+        # self.mapping_network = MappingNetwork(z_dim=hidden_dim, w_dim=hidden_dim, num_layers=num_layers)
+
+    def transform_rel(self, token_traj, prev_pos, prev_heading=None):
+        if prev_heading is None:
+            diff_xy = prev_pos[:, :, -1, :] - prev_pos[:, :, -2, :]
+            prev_heading = torch.arctan2(diff_xy[:, :, 1], diff_xy[:, :, 0])
+
+        num_agent, num_step, traj_num, traj_dim = token_traj.shape
+        cos, sin = prev_heading.cos(), prev_heading.sin()
+        rot_mat = torch.zeros((num_agent, num_step, 2, 2), device=prev_heading.device)
+        rot_mat[:, :, 0, 0] = cos
+        rot_mat[:, :, 0, 1] = -sin
+        rot_mat[:, :, 1, 0] = sin
+        rot_mat[:, :, 1, 1] = cos
+        agent_diff_rel = torch.bmm(token_traj.view(-1, traj_num, 2), rot_mat.view(-1, 2, 2)).view(num_agent, num_step, traj_num, traj_dim)
+        agent_pred_rel = agent_diff_rel + prev_pos[:, :, -1:, :]
+        return agent_pred_rel
+
+    def agent_token_embedding(self, data, agent_token_index, agent_state, pos_a, head_a, inference=False,
+                              filter_mask=None, av_index=None):
+
+        if filter_mask is None:
+            filter_mask = torch.ones_like(agent_state[:, 2], dtype=torch.bool)
+
+        num_agent, num_step, traj_dim = pos_a.shape # traj_dim=2
+        agent_type = data['agent']['type'][filter_mask]
+        veh_mask = (agent_type == 0)
+        ped_mask = (agent_type == 1)
+        cyc_mask = (agent_type == 2)
+
+        # set the position of invalid agents to the position of ego agent
+        # note here we only set invalid steps BEFORE the bos token!
+        # is_invalid = agent_state == self.invalid_state
+        # is_bos = agent_state == self.enter_state
+        # bos_index = torch.where(is_bos.any(dim=1), torch.argmax(is_bos.long(), dim=1), torch.tensor(0))
+        # bos_mask = torch.arange(num_step).expand(num_agent, -1).to(agent_state.device) < bos_index[:, None]
+        # is_invalid[~bos_mask] = False
+
+        # ego_pos_a = pos_a[av_index].clone()
+        # ego_head_vector_a = head_vector_a[av_index].clone()
+        # pos_a[is_invalid] = ego_pos_a[None, :].repeat(pos_a.shape[0], 1, 1)[is_invalid]
+        # head_vector_a[is_invalid] = ego_head_vector_a[None, :].repeat(head_vector_a.shape[0], 1, 1)[is_invalid]
+
+        motion_vector_a, head_vector_a = self.build_vector_a(pos_a, head_a, agent_state)
+
+        trajectory_token_veh = torch.from_numpy(self.trajectory_token['veh']).clone().to(pos_a.device).to(torch.float)
+        trajectory_token_ped = torch.from_numpy(self.trajectory_token['ped']).clone().to(pos_a.device).to(torch.float)
+        trajectory_token_cyc = torch.from_numpy(self.trajectory_token['cyc']).clone().to(pos_a.device).to(torch.float)
+        self.agent_token_emb_veh = self.token_emb_veh(trajectory_token_veh.view(trajectory_token_veh.shape[0], -1)) # (token_size, 8)
+        self.agent_token_emb_ped = self.token_emb_ped(trajectory_token_ped.view(trajectory_token_ped.shape[0], -1))
+        self.agent_token_emb_cyc = self.token_emb_cyc(trajectory_token_cyc.view(trajectory_token_cyc.shape[0], -1))
+
+        # add bos token embedding
+        self.agent_token_emb_veh = torch.cat([self.agent_token_emb_veh, self.bos_token_emb(torch.zeros(1, device=pos_a.device).long())])
+        self.agent_token_emb_ped = torch.cat([self.agent_token_emb_ped, self.bos_token_emb(torch.zeros(1, device=pos_a.device).long())])
+        self.agent_token_emb_cyc = torch.cat([self.agent_token_emb_cyc, self.bos_token_emb(torch.zeros(1, device=pos_a.device).long())])
+
+        # add invalid token embedding
+        self.agent_token_emb_veh = torch.cat([self.agent_token_emb_veh, self.no_token_emb(torch.zeros(1, device=pos_a.device).long())])
+        self.agent_token_emb_ped = torch.cat([self.agent_token_emb_ped, self.no_token_emb(torch.zeros(1, device=pos_a.device).long())])
+        self.agent_token_emb_cyc = torch.cat([self.agent_token_emb_cyc, self.no_token_emb(torch.zeros(1, device=pos_a.device).long())])
+
+        if inference:
+            agent_token_traj_all = torch.zeros((num_agent, self.token_size, self.shift + 1, 4, 2), device=pos_a.device)
+            trajectory_token_all_veh = torch.from_numpy(self.trajectory_token_all['veh']).clone().to(pos_a.device).to(torch.float)
+            trajectory_token_all_ped = torch.from_numpy(self.trajectory_token_all['ped']).clone().to(pos_a.device).to(torch.float)
+            trajectory_token_all_cyc = torch.from_numpy(self.trajectory_token_all['cyc']).clone().to(pos_a.device).to(torch.float)
+            agent_token_traj_all[veh_mask] = torch.cat(
+                [trajectory_token_all_veh[:, :self.shift], trajectory_token_veh[:, None, ...]], dim=1)
+            agent_token_traj_all[ped_mask] = torch.cat(
+                [trajectory_token_all_ped[:, :self.shift], trajectory_token_ped[:, None, ...]], dim=1)
+            agent_token_traj_all[cyc_mask] = torch.cat(
+                [trajectory_token_all_cyc[:, :self.shift], trajectory_token_cyc[:, None, ...]], dim=1)
+
+        # additional token embeddings are already added -> -1: invalid, -2: bos
+        agent_token_emb = torch.zeros((num_agent, num_step, self.hidden_dim), device=pos_a.device)
+        agent_token_emb[veh_mask] = self.agent_token_emb_veh[agent_token_index[veh_mask]]
+        agent_token_emb[ped_mask] = self.agent_token_emb_ped[agent_token_index[ped_mask]]
+        agent_token_emb[cyc_mask] = self.agent_token_emb_cyc[agent_token_index[cyc_mask]]
+
+        # 'vehicle', 'pedestrian', 'cyclist', 'background'
+        is_invalid = (agent_state == self.invalid_state) & (agent_state != self.enter_state)
+        agent_types = data['agent']['type'][filter_mask].long().repeat_interleave(repeats=num_step, dim=0)
+        agent_types[is_invalid.reshape(-1)] = self.all_agent_type.index('invalid')
+        agent_shapes = data['agent']['shape'][filter_mask, self.num_historical_steps - 1, :].repeat_interleave(repeats=num_step, dim=0)
+        agent_shapes[is_invalid.reshape(-1)] = self.invalid_shape_value
+
+        categorical_embs = [self.type_a_emb(agent_types), self.shape_emb(agent_shapes)]
+        feature_a = torch.stack(
+            [torch.norm(motion_vector_a[:, :, :2], p=2, dim=-1),
+             angle_between_2d_vectors(ctr_vector=head_vector_a, nbr_vector=motion_vector_a[:, :, :2]),
+             ], dim=-1) # (num_agent, num_shifted_step, 2)
+
+        x_a = self.x_a_emb(continuous_inputs=feature_a.view(-1, feature_a.size(-1)),
+                           categorical_embs=categorical_embs)
+        x_a = x_a.view(-1, num_step, self.hidden_dim) # (num_agent, num_step, hidden_dim)
+
+        s_a = self.state_a_emb(agent_state.reshape(-1).long()).reshape(num_agent, num_step, self.hidden_dim)
+        feat_a = torch.cat((agent_token_emb, x_a, s_a), dim=-1) # (num_agent, num_step, hidden_dim * 3)
+        feat_a = self.fusion_emb(feat_a) # (num_agent, num_step, hidden_dim)
+
+        # seed agent feature
+        motion_vector_seed = motion_vector_a[av_index : av_index + 1]
+        head_vector_seed = head_vector_a[av_index : av_index + 1]
+        feat_seed = self.build_invalid_agent_feature(num_step, pos_a.device, type_index=self.all_agent_type.index('seed'),
+                                                     motion_vector=motion_vector_seed, head_vector=head_vector_seed)
+
+        # replace the features of steps before bos of valid agents with the corresponding invalid agent features
+        # is_bos = agent_state == self.enter_state
+        # is_eos = agent_state == self.exit_state
+        # bos_index = torch.where(is_bos.any(dim=1), torch.argmax(is_bos.long(), dim=1), torch.tensor(0))
+        # eos_index = torch.where(is_eos.any(dim=1), torch.argmax(is_eos.long(), dim=1), torch.tensor(num_step - 1))
+        # is_before_bos = torch.arange(num_step).expand(num_agent, -1).to(agent_state.device) < bos_index[:, None]
+        # is_after_eos = torch.arange(num_step).expand(num_agent, -1).to(agent_state.device) > eos_index[:, None] + 1
+        # feat_ina = self.build_invalid_agent_feature(num_step, pos_a.device)
+        # feat_a[is_before_bos | is_after_eos] = feat_ina.repeat(num_agent, 1, 1)[is_before_bos | is_after_eos]
+
+        # print("train")
+        # is_bos = agent_state == self.enter_state
+        # is_eos = agent_state == self.exit_state
+        # bos_index = torch.where(is_bos.any(dim=1), torch.argmax(is_bos.long(), dim=1), torch.tensor(0))
+        # eos_index = torch.where(is_eos.any(dim=1), torch.argmax(is_eos.long(), dim=1), torch.tensor(num_step - 1))
+        # mask = torch.arange(num_step).expand(num_agent, -1).to(agent_state.device)
+        # mask = (mask >= bos_index[:, None]) & (mask <= eos_index[:, None] + 1)
+        # is_invalid[mask] = False
+        # print(feat_a.sum(dim=-1)[is_invalid])
+
+        feat_a = torch.cat([feat_a, feat_seed], dim=0) # (num_agent + 1, num_step, hidden_dim)
+
+        # feat_a_sum = feat_a.sum(dim=-1)
+        # for a in range(num_agent):
+        #     print(f"agent {a}:")
+        #     print(f"state: {agent_state[a, :]}")
+        #     print(f"feat_a_sum: {feat_a_sum[a, :]}")
+        # exit(1)
+
+        if inference:
+            return feat_a, head_vector_a, agent_token_traj_all, agent_token_emb, categorical_embs
+        else:
+            return feat_a, head_vector_a
+
+    def build_vector_a(self, pos_a, head_a, state_a):
+        num_agent = pos_a.shape[0]
+
+        motion_vector_a = torch.cat([pos_a.new_zeros(num_agent, 1, self.input_dim),
+                                     pos_a[:, 1:] - pos_a[:, :-1]], dim=1)
+
+        # update the relative motion/head vectors
+        is_bos = state_a == self.enter_state
+        motion_vector_a[is_bos] = self.motion_gap
+
+        is_last_eos = state_a.roll(shifts=1, dims=1) == self.exit_state
+        is_last_eos[:, 0] = False
+        motion_vector_a[is_last_eos] = -self.motion_gap
+
+        head_vector_a = torch.stack([head_a.cos(), head_a.sin()], dim=-1)
+
+        return motion_vector_a, head_vector_a
+
+    def build_invalid_agent_feature(self, num_step, device, motion_vector=None, head_vector=None, type_index=None, shape_value=None):
+        invalid_agent_token_emb = self.no_token_emb(torch.zeros(1, device=device).long())[:, None].repeat(1, num_step, 1)
+
+        if motion_vector is None or head_vector is None:
+            motion_vector = torch.zeros((1, num_step, 2), device=device)
+            head_vector = torch.stack([torch.cos(torch.zeros(1, device=device)), torch.sin(torch.zeros(1, device=device))], dim=-1)[:, None, :].repeat(1, num_step, 1)
+
+        feature_ina = torch.stack(
+            [torch.norm(motion_vector[:, :, :2], p=2, dim=-1),
+             angle_between_2d_vectors(ctr_vector=head_vector, nbr_vector=motion_vector[:, :, :2]),
+             ], dim=-1)
+
+        if type_index is None:
+            type_index = self.all_agent_type.index('invalid')
+        if shape_value is None:
+            shape_value = torch.full((1, 3), self.invalid_shape_value, device=device)
+
+        categorical_embs_ina = [self.type_a_emb(torch.tensor([type_index], device=device)),
+                                 self.shape_emb(shape_value)]
+        x_ina = self.x_a_emb(continuous_inputs=feature_ina.view(-1, feature_ina.size(-1)),
+                             categorical_embs=categorical_embs_ina)
+        x_ina = x_ina.view(-1, num_step, self.hidden_dim) # (1, num_step, hidden_dim)
+
+        s_ina = self.state_a_emb(torch.tensor([self.invalid_state], device=device))[:, None].repeat(1, num_step, 1) # NOTE: do not use `expand`
+
+        feat_ina = torch.cat((invalid_agent_token_emb, x_ina, s_ina), dim=-1)
+        feat_ina = self.fusion_emb(feat_ina) # (1, num_step, hidden_dim)
+
+        return feat_ina
+
+    def build_temporal_edge(self, pos_a, head_a, head_vector_a, state_a, mask, inference_mask=None, av_index=None):
+
+        num_agent = pos_a.shape[0]
+        hist_mask = mask.clone()
+
+        if not self.temporal_attn_to_invalid:
+            hist_mask[state_a == self.invalid_state] = False
+
+        # set the position of invalid agents to the position of ego agent
+        ego_pos_a = pos_a[av_index].clone() # (num_step, 2)
+        ego_head_a = head_a[av_index].clone()
+        ego_head_vector_a = head_vector_a[av_index].clone()
+        ego_state_a = state_a[av_index].clone()
+        # is_invalid = state_a == self.invalid_state
+        # pos_a[is_invalid] = ego_pos_a[None, :].repeat(pos_a.shape[0], 1, 1)[is_invalid]
+        # head_a[is_invalid] = ego_head_a[None, :].repeat(head_a.shape[0], 1)[is_invalid]
+
+        # add seed agent
+        pos_a = torch.cat([pos_a, ego_pos_a[None]], dim=0)
+        head_a = torch.cat([head_a, ego_head_a[None]], dim=0)
+        state_a = torch.cat([state_a, ego_state_a[None]], dim=0)
+        head_vector_a = torch.cat([head_vector_a, ego_head_vector_a[None]], dim=0)
+        hist_mask = torch.cat([hist_mask, torch.ones_like(hist_mask[0:1])], dim=0).bool()
+        if not self.temporal_attn_seed:
+            hist_mask[-1:] = False
+            if inference_mask is not None:
+                inference_mask[-1:] = False
+
+        pos_t = pos_a.reshape(-1, self.input_dim) # (num_agent * num_step, ...)
+        head_t = head_a.reshape(-1)
+        head_vector_t = head_vector_a.reshape(-1, 2)
+
+        # for those invalid agents won't predict any motion token, we don't attend to them
+        is_bos = state_a == self.enter_state
+        is_bos[-1] = False
+        bos_index = torch.where(is_bos.any(dim=1), torch.argmax(is_bos.long(), dim=1), torch.tensor(0))
+        motion_predict_start_index = torch.clamp(bos_index - self.time_span / self.shift + 1, min=0)
+        motion_predict_mask = torch.arange(hist_mask.shape[1]).expand(hist_mask.shape[0], -1).to(hist_mask.device)
+        motion_predict_mask = motion_predict_mask >= motion_predict_start_index[:, None]
+        hist_mask[~motion_predict_mask] = False
+
+        if self.hist_mask and self.training:
+            hist_mask[
+                torch.arange(mask.shape[0]).unsqueeze(1), torch.randint(0, mask.shape[1], (num_agent, 10))] = False
+            mask_t = hist_mask.unsqueeze(2) & hist_mask.unsqueeze(1)
+        elif inference_mask is not None:
+            mask_t = hist_mask.unsqueeze(2) & inference_mask.unsqueeze(1)
+        else:
+            mask_t = hist_mask.unsqueeze(2) & hist_mask.unsqueeze(1)
+
+        # mask_t: (num_agent, 18, 18), edge_index_t: (2, num_edge)
+        edge_index_t = dense_to_sparse(mask_t)[0]
+        edge_index_t = edge_index_t[:, (edge_index_t[1] - edge_index_t[0] > 0) &
+                                       (edge_index_t[1] - edge_index_t[0] <= self.time_span / self.shift)]
+        rel_pos_t = pos_t[edge_index_t[0]] - pos_t[edge_index_t[1]]
+        rel_head_t = wrap_angle(head_t[edge_index_t[0]] - head_t[edge_index_t[1]])
+
+        # FIXME relative motion/head for bos/eos token
+        # is_next_bos = state_a.roll(shifts=-1, dims=1) == self.enter_state
+        # is_next_bos[:, -1] = False # the last step
+        # is_next_bos_t = is_next_bos.reshape(-1)
+        # rel_pos_t[is_next_bos_t[edge_index_t[0]]] = -self.bos_motion
+        # rel_pos_t[is_next_bos_t[edge_index_t[1]]] = self.bos_motion
+        # rel_head_t[is_next_bos_t[edge_index_t[0]]] = -torch.pi
+        # rel_head_t[is_next_bos_t[edge_index_t[1]]] = torch.pi
+
+        # is_last_eos = state_a.roll(shifts=1, dims=1) == self.exit_state
+        # is_last_eos[:, 0] = False # the first step
+        # is_last_eos_t = is_last_eos.reshape(-1)
+        # rel_pos_t[is_last_eos_t[edge_index_t[0]]] = -self.bos_motion
+        # rel_pos_t[is_last_eos_t[edge_index_t[1]]] = self.bos_motion
+        # rel_head_t[is_last_eos_t[edge_index_t[0]]] = -torch.pi
+        # rel_head_t[is_last_eos_t[edge_index_t[1]]] = torch.pi
+
+        # handle the bos token of ego agent
+        # is_invalid = state_a == self.invalid_state
+        # is_invalid_t = is_invalid.reshape(-1)
+        # is_ego_bos = (ego_state_a == self.enter_state)[None, :].expand(num_agent + 1, -1)
+        # is_ego_bos_t = is_ego_bos.reshape(-1)
+        # rel_pos_t[is_invalid_t[edge_index_t[0]] & is_ego_bos_t[edge_index_t[0]]] = 0.
+        # rel_pos_t[is_invalid_t[edge_index_t[1]] & is_ego_bos_t[edge_index_t[1]]] = 0.
+        # rel_head_t[is_invalid_t[edge_index_t[0]] & is_ego_bos_t[edge_index_t[0]]] = 0.
+        # rel_head_t[is_invalid_t[edge_index_t[1]] & is_ego_bos_t[edge_index_t[1]]] = 0.
+
+        # handle the invalid steps
+        is_invalid = state_a == self.invalid_state
+        is_invalid_t = is_invalid.reshape(-1)
+        rel_pos_t[is_invalid_t[edge_index_t[0]]] = -self.motion_gap
+        rel_pos_t[is_invalid_t[edge_index_t[1]]] = self.motion_gap
+        rel_head_t[is_invalid_t[edge_index_t[0]]] = -self.heading_gap
+        rel_head_t[is_invalid_t[edge_index_t[1]]] = self.heading_gap
+
+        r_t = torch.stack(
+            [torch.norm(rel_pos_t[:, :2], p=2, dim=-1),
+             angle_between_2d_vectors(ctr_vector=head_vector_t[edge_index_t[1]], nbr_vector=rel_pos_t[:, :2]),
+             rel_head_t,
+             edge_index_t[0] - edge_index_t[1]], dim=-1)
+        r_t = self.r_t_emb(continuous_inputs=r_t, categorical_embs=None)
+
+        return edge_index_t, r_t
+
+    def build_interaction_edge(self, pos_a, head_a, head_vector_a, state_a, batch_s, mask_a, inference_mask=None, av_index=None):
+        num_agent, num_step, _ = pos_a.shape
+
+        pos_a = torch.cat([pos_a, pos_a[av_index][None]], dim=0)
+        head_a = torch.cat([head_a, head_a[av_index][None]], dim=0)
+        state_a = torch.cat([state_a, state_a[av_index][None]], dim=0)
+        head_vector_a = torch.cat([head_vector_a, head_vector_a[av_index][None]], dim=0)
+
+        pos_s = pos_a.transpose(0, 1).reshape(-1, self.input_dim)
+        head_s = head_a.transpose(0, 1).reshape(-1)
+        head_vector_s = head_vector_a.transpose(0, 1).reshape(-1, 2)
+        if inference_mask is not None:
+            mask_a = mask_a & inference_mask
+        mask_s = mask_a.transpose(0, 1).reshape(-1)
+
+        # seed agent
+        mask_seed = state_a[av_index] != self.invalid_state
+        pos_seed = pos_a[av_index]
+        edge_index_seed2a = radius(x=pos_seed[:, :2], y=pos_s[:, :2], r=self.pl2seed_radius,
+                                   batch_x=torch.arange(num_step).to(pos_s.device), batch_y=batch_s, max_num_neighbors=300)
+        edge_index_seed2a = edge_index_seed2a[:, mask_s[edge_index_seed2a[0]] & mask_seed[edge_index_seed2a[1]]]
+
+        # convert to global index (must be unilateral connection)
+        edge_index_seed2a[1, :] = (edge_index_seed2a[1, :] + 1) * (num_agent + 1) - 1
+
+        # build agent2agent bilateral connection
+        edge_index_a2a = radius_graph(x=pos_s[:, :2], r=self.a2a_radius, batch=batch_s, loop=False,
+                                      max_num_neighbors=300)
+        edge_index_a2a = subgraph(subset=mask_s, edge_index=edge_index_a2a)[0]
+
+        # add the edges which connect seed agents
+        edge_index_a2a = torch.cat([edge_index_a2a, edge_index_seed2a], dim=-1)
+
+        # set the position of invalid agents to the position of ego agent
+        # ego_pos_a = pos_a[av_index].clone()
+        # ego_head_a = head_a[av_index].clone()
+        # is_invalid = state_a == self.invalid_state
+        # pos_a[is_invalid] = ego_pos_a[None, :].repeat(pos_a.shape[0], 1, 1)[is_invalid]
+        # head_a[is_invalid] = ego_head_a[None, :].repeat(head_a.shape[0], 1)[is_invalid]
+
+        rel_pos_a2a = pos_s[edge_index_a2a[0]] - pos_s[edge_index_a2a[1]]
+        rel_head_a2a = wrap_angle(head_s[edge_index_a2a[0]] - head_s[edge_index_a2a[1]])
+
+        # relative motion/head for bos/eos token
+        # is_bos = state_a == self.enter_state
+        # is_bos_s = is_bos.transpose(0, 1).reshape(-1)
+        # rel_pos_a2a[is_bos_s[edge_index_a2a[0]]] = -self.bos_motion
+        # rel_pos_a2a[is_bos_s[edge_index_a2a[1]]] = self.bos_motion
+        # rel_head_a2a[is_bos_s[edge_index_a2a[0]]] = -torch.pi
+        # rel_head_a2a[is_bos_s[edge_index_a2a[1]]] = torch.pi
+
+        # is_last_eos = state_a.roll(shifts=-1, dims=1) == self.exit_state
+        # is_last_eos[:, 0] = False # first step
+        # is_last_eos_s = is_last_eos.transpose(0, 1).reshape(-1)
+        # rel_pos_a2a[is_last_eos_s[edge_index_a2a[0]]] = -self.bos_motion
+        # rel_pos_a2a[is_last_eos_s[edge_index_a2a[1]]] = self.bos_motion
+        # rel_head_a2a[is_last_eos_s[edge_index_a2a[0]]] = -torch.pi
+        # rel_head_a2a[is_last_eos_s[edge_index_a2a[1]]] = torch.pi
+
+        # handle the bos token of ego agent
+        # is_invalid = state_a == self.invalid_state
+        # is_invalid_s = is_invalid.transpose(0, 1).reshape(-1)
+        # is_ego_bos = (state_a[av_index] == self.enter_state)[None, :].expand(num_agent + 1, -1)
+        # is_ego_bos_s = is_ego_bos.transpose(0, 1).reshape(-1)
+        # rel_pos_a2a[is_invalid_s[edge_index_a2a[0]] & is_ego_bos_s[edge_index_a2a[0]]] = 0.
+        # rel_pos_a2a[is_invalid_s[edge_index_a2a[1]] & is_ego_bos_s[edge_index_a2a[1]]] = 0.
+        # rel_head_a2a[is_invalid_s[edge_index_a2a[0]] & is_ego_bos_s[edge_index_a2a[0]]] = 0.
+        # rel_head_a2a[is_invalid_s[edge_index_a2a[1]] & is_ego_bos_s[edge_index_a2a[1]]] = 0.
+
+        # handle the invalid steps
+        is_invalid = state_a == self.invalid_state
+        is_invalid_s = is_invalid.transpose(0, 1).reshape(-1)
+        rel_pos_a2a[is_invalid_s[edge_index_a2a[0]]] = -self.motion_gap
+        rel_pos_a2a[is_invalid_s[edge_index_a2a[1]]] = self.motion_gap
+        rel_head_a2a[is_invalid_s[edge_index_a2a[0]]] = -self.heading_gap
+        rel_head_a2a[is_invalid_s[edge_index_a2a[1]]] = self.heading_gap
+
+        r_a2a = torch.stack(
+            [torch.norm(rel_pos_a2a[:, :2], p=2, dim=-1),
+             angle_between_2d_vectors(ctr_vector=head_vector_s[edge_index_a2a[1]], nbr_vector=rel_pos_a2a[:, :2]),
+             rel_head_a2a], dim=-1)
+        r_a2a = self.r_a2a_emb(continuous_inputs=r_a2a, categorical_embs=None)
+
+        return edge_index_a2a, r_a2a
+
+    def build_map2agent_edge(self, data, num_step, pos_a, head_a, head_vector_a, state_a, batch_s, batch_pl,
+                             mask, inference_mask=None, av_index=None):
+
+        num_agent, num_step, _ = pos_a.shape
+
+        mask_pl2a = mask.clone()
+        if inference_mask is not None:
+            mask_pl2a = mask_pl2a & inference_mask
+        mask_pl2a = mask_pl2a.transpose(0, 1).reshape(-1)
+
+        pos_a = torch.cat([pos_a, pos_a[av_index][None]], dim=0)
+        state_a = torch.cat([state_a, state_a[av_index][None]], dim=0)
+        head_a = torch.cat([head_a, head_a[av_index][None]], dim=0)
+        head_vector_a = torch.cat([head_vector_a, head_vector_a[av_index][None]], dim=0)
+
+        pos_s = pos_a.transpose(0, 1).reshape(-1, self.input_dim)
+        head_s = head_a.transpose(0, 1).reshape(-1)
+        head_vector_s = head_vector_a.transpose(0, 1).reshape(-1, 2)
+
+        ori_pos_pl = data['pt_token']['position'][:, :self.input_dim].contiguous()
+        ori_orient_pl = data['pt_token']['orientation'].contiguous()
+        pos_pl = ori_pos_pl.repeat(num_step, 1) # not `repeat_interleave`
+        orient_pl = ori_orient_pl.repeat(num_step)
+
+        # seed agent
+        mask_seed = state_a[av_index] != self.invalid_state
+        pos_seed = pos_a[av_index]
+        edge_index_pl2seed = radius(x=pos_seed[:, :2], y=pos_pl[:, :2], r=self.pl2seed_radius,
+                                   batch_x=torch.arange(num_step).to(pos_s.device), batch_y=batch_pl, max_num_neighbors=600)
+        edge_index_pl2seed = edge_index_pl2seed[:, mask_seed[edge_index_pl2seed[1]]]
+
+        # convert to global index
+        edge_index_pl2seed[1, :] = (edge_index_pl2seed[1, :] + 1) * (num_agent + 1) - 1
+
+        # build map2agent directed graph
+        # edge_index_pl2a[0]: pl token; edge_index_pl2a[1]: agent token
+        edge_index_pl2a = radius(x=pos_s[:, :2], y=pos_pl[:, :2], r=self.pl2a_radius,
+                                 batch_x=batch_s, batch_y=batch_pl, max_num_neighbors=300)
+        # We force invalid agents to interact with **all** (visible in current window) map tokens
+        # invalid_node_index_a = torch.where(bos_state_s.bool())[0]
+        # sampled_node_index_m = torch.arange(ori_pos_pl.shape[0]).to(pos_pl.device)
+        # if kwargs.get('sample_pt_indices', None) is not None:
+        #     sampled_node_index_m = sampled_node_index_m[kwargs['sample_pt_indices'].long()] 
+        # grid_a, grid_b = torch.meshgrid(sampled_node_index_m, invalid_node_index_a, indexing='ij')
+        # invalid_edge_index_pl2a = torch.stack([grid_a.reshape(-1), grid_b.reshape(-1)], dim=0)
+        # edge_index_pl2a = torch.concat([edge_index_pl2a, invalid_edge_index_pl2a], dim=-1)
+        # remove the edges which connect with motion-invalid agents
+        edge_index_pl2a = edge_index_pl2a[:, mask_pl2a[edge_index_pl2a[1]]]
+
+        # add the edges which connect seed agents with map tokens
+        edge_index_pl2a = torch.cat([edge_index_pl2a, edge_index_pl2seed], dim=-1)
+
+        # set the position of invalid agents to the position of ego agent
+        # ego_pos_a = pos_a[av_index].clone()
+        # ego_head_a = head_a[av_index].clone()
+        # is_invalid = state_a == self.invalid_state
+        # pos_a[is_invalid] = ego_pos_a[None, :].repeat(pos_a.shape[0], 1, 1)[is_invalid]
+        # head_a[is_invalid] = ego_head_a[None, :].repeat(head_a.shape[0], 1)[is_invalid]
+
+        rel_pos_pl2a = pos_pl[edge_index_pl2a[0]] - pos_s[edge_index_pl2a[1]]
+        rel_orient_pl2a = wrap_angle(orient_pl[edge_index_pl2a[0]] - head_s[edge_index_pl2a[1]])
+
+        # handle the invalid steps
+        is_invalid = state_a == self.invalid_state
+        is_invalid_s = is_invalid.transpose(0, 1).reshape(-1)
+        rel_pos_pl2a[is_invalid_s[edge_index_pl2a[1]]] = self.motion_gap
+        rel_orient_pl2a[is_invalid_s[edge_index_pl2a[1]]] = self.heading_gap
+
+        r_pl2a = torch.stack(
+            [torch.norm(rel_pos_pl2a[:, :2], p=2, dim=-1),
+             angle_between_2d_vectors(ctr_vector=head_vector_s[edge_index_pl2a[1]], nbr_vector=rel_pos_pl2a[:, :2]),
+             rel_orient_pl2a], dim=-1)
+        r_pl2a = self.r_pt2a_emb(continuous_inputs=r_pl2a, categorical_embs=None)
+
+        return edge_index_pl2a, r_pl2a
+
+    def get_inputs(self, data: HeteroData) -> Dict[str, torch.Tensor]:
+
+        pos_a = data['agent']['token_pos']
+        head_a = data['agent']['token_heading']
+        agent_category = data['agent']['category']
+        agent_token_index = data['agent']['token_idx']
+        agent_state_index = data['agent']['state_idx']
+        mask = data['agent']['raw_agent_valid_mask'].clone()
+        # mask[agent_category != 3] = False
+
+        if not self.predict_state:
+            agent_state_index = None
+
+        next_token_index_gt = agent_token_index.roll(shifts=-1, dims=1)
+        next_state_index_gt = agent_state_index.roll(shifts=-1, dims=1)
+
+        if self.predict_state:
+            next_token_eval_mask = mask.clone()
+            next_token_eval_mask = next_token_eval_mask * next_token_eval_mask.roll(shifts=1, dims=1)
+            bos_token_index = torch.nonzero(agent_state_index == 2)
+            eos_token_index = torch.nonzero(agent_state_index == 3)
+            next_token_eval_mask[bos_token_index[:, 0], bos_token_index[:, 1]] = 1
+            for eos_token_index_ in eos_token_index:
+                if not next_token_eval_mask[eos_token_index_[0], eos_token_index_[1]]:
+                    next_token_eval_mask[eos_token_index_[0], eos_token_index_[1]:] = 0
+            next_token_eval_mask = next_token_eval_mask.roll(shifts=-1, dims=1)
+            # TODO: next_state_eval_mask !!!
+
+            if next_token_index_gt[next_token_eval_mask].min() < 0:
+                raise RuntimeError()
+
+        next_token_eval_mask[:, -1] = False
+
+        return {'token_pos': pos_a,
+                'token_heading': head_a,
+                'agent_category': agent_category,
+                'next_token_idx_gt': next_token_index_gt,
+                'next_state_idx_gt': next_state_index_gt,
+                'next_token_eval_mask': next_token_eval_mask,
+                'raw_agent_valid_mask': data['agent']['raw_agent_valid_mask'],
+                }
+
+    def forward(self,
+                data: HeteroData,
+                map_enc: Mapping[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
+
+        pos_a = data['agent']['token_pos'].clone() # (num_agent, num_shifted_step, 2)
+        head_a = data['agent']['token_heading'].clone() # (num_agent, num_shifted_step)
+        num_agent, num_step, traj_dim = pos_a.shape # e.g. (50, 18, 2)
+        agent_category = data['agent']['category'].clone() # (num_agent,)
+        agent_token_index = data['agent']['token_idx'].clone() # (num_agent, num_step)
+        agent_state_index = data['agent']['state_idx'].clone() # (num_agent, num_step)
+        agent_type_index = data['agent']['type'].clone() # (num_agent, num_step)
+        agent_enter_pl_token_idx = None
+        agent_enter_offset_token_idx = None
+
+        device = pos_a.device
+
+        seed_step_mask = agent_state_index[:, 1:] == self.enter_state
+        if torch.any(seed_step_mask.sum(dim=0) > self.seed_size):
+            print(agent_state_index)
+            print(agent_state_index.shape)
+            print(seed_step_mask.long())
+            print(seed_step_mask.sum(dim=0))
+            raise RuntimeError(f"Seed size {self.seed_size} is too small.")
+
+        # fix pos and head of invalid agents
+        av_index = int(data['agent']['av_index'])
+        # ego_pos_a = pos_a[av_index].clone() # (num_shifted_step, 2)
+        # ego_head_vector_a = head_vector_a[av_index] # (num_shifted_step, 2)
+        # is_invalid = agent_state_index == self.invalid_state
+        # pos_a[is_invalid] = ego_pos_a[None, :].expand(pos_a.shape[0], -1, -1)[is_invalid]
+        # head_vector_a[is_invalid] = ego_head_vector_a[None, :].expand(head_vector_a.shape[0], -1, -1)[is_invalid]
+
+        if not self.predict_state:
+            agent_state_index = None
+
+        feat_a, head_vector_a = self.agent_token_embedding(data, agent_token_index, agent_state_index, pos_a, head_a, av_index=av_index)
+
+        # build masks
+        mask = data['agent']['raw_agent_valid_mask'].clone()
+        temporal_mask = mask.clone()
+        interact_mask = mask.clone()
+        if self.predict_state:
+
+            agent_enter_offset_token_idx = data['agent']['neighbor_token_idx']
+            agent_enter_pl_token_idx = data['agent']['map_bos_token_idx']
+            agent_enter_pl_token_id = data['agent']['map_bos_token_id']
+
+            is_bos = agent_state_index == self.enter_state
+            is_eos = agent_state_index == self.exit_state
+            bos_index = torch.where(is_bos.any(dim=1), torch.argmax(is_bos.long(), dim=1), torch.tensor(0))
+            eos_index = torch.where(is_eos.any(dim=1), torch.argmax(is_eos.long(), dim=1), torch.tensor(num_step - 1)) # not `-1`
+
+            temporal_mask = torch.ones_like(mask)
+            motion_mask = torch.arange(mask.shape[1]).expand(mask.shape[0], -1).to(device)
+            motion_mask = (motion_mask > bos_index[:, None]) & (motion_mask <= eos_index[:, None])
+            temporal_mask[motion_mask] = mask[motion_mask]
+
+            interact_mask[agent_state_index == self.enter_state] = True
+            interact_mask = torch.cat([interact_mask, torch.ones_like(interact_mask[:1])]).bool() # placeholder
+
+        edge_index_t, r_t = self.build_temporal_edge(pos_a, head_a, head_vector_a, agent_state_index, temporal_mask,
+                                                     av_index=av_index)
+
+        # +1: placeholder for seed agent
+        # if isinstance(data, Batch):
+        #     print(data['agent']['batch'], data.num_graphs)
+        #     batch_s = torch.cat([data['agent']['batch'] + data.num_graphs * t for t in range(num_step)], dim=0)
+        #     batch_pl = torch.cat([data['pt_token']['batch'] + data.num_graphs * t for t in range(num_step)], dim=0)
+        # else:
+        batch_s = torch.arange(num_step, device=device).repeat_interleave(data['agent']['num_nodes'] + 1)
+        batch_pl = torch.arange(num_step, device=device).repeat_interleave(data['pt_token']['num_nodes'])
+
+        edge_index_a2a, r_a2a = self.build_interaction_edge(pos_a, head_a, head_vector_a, agent_state_index, batch_s,
+                                                            interact_mask, av_index=av_index)
+
+        agent_category = torch.cat([agent_category, torch.full(agent_category[-1:].shape, 3, device=device)])
+        interact_mask[agent_category != 3] = False
+        edge_index_pl2a, r_pl2a = self.build_map2agent_edge(data, num_step, pos_a, head_a, head_vector_a,
+                                                            agent_state_index, batch_s, batch_pl, interact_mask, av_index=av_index)
+
+        # mapping network
+        # z = torch.randn(num_agent, self.hidden_dim).to(feat_a.device)
+        # w = self.mapping_network(z)
+
+        for i in range(self.num_layers):
+
+            # feat_a = feat_a + w[:, None]
+
+            feat_a = feat_a.reshape(-1, self.hidden_dim) # (num_agent, num_step, hidden_dim) -> (seq_len, hidden_dim)
+            feat_a = self.t_attn_layers[i](feat_a, r_t, edge_index_t)
+
+            feat_a = feat_a.reshape(-1, num_step, self.hidden_dim).transpose(0, 1).reshape(-1, self.hidden_dim)
+            feat_a = self.pt2a_attn_layers[i]((
+                map_enc['x_pt'].repeat_interleave(repeats=num_step, dim=0).reshape(-1, num_step, self.hidden_dim).transpose(0, 1).reshape(
+                -1, self.hidden_dim), feat_a), r_pl2a, edge_index_pl2a)
+
+            feat_a = self.a2a_attn_layers[i](feat_a, r_a2a, edge_index_a2a)
+            feat_a = feat_a.reshape(num_step, -1, self.hidden_dim).transpose(0, 1)
+
+        # next motion token
+        next_token_prob = self.token_predict_head(feat_a[:-1]) # (num_agent, num_step, token_size)
+        next_token_prob_softmax = torch.softmax(next_token_prob, dim=-1)
+        _, next_token_idx = torch.topk(next_token_prob_softmax, k=10, dim=-1) # (num_agent, num_step, 10)
+
+        next_token_index_gt = agent_token_index.roll(shifts=-1, dims=1)
+
+        # next state token
+        next_state_prob = self.state_predict_head(feat_a[:-1])
+        next_state_idx = next_state_prob.softmax(dim=-1).argmax(dim=-1, keepdim=True) # (num_agent, num_step, 1)
+
+        next_state_index_gt = agent_state_index.roll(shifts=-1, dims=1) # (invalid, valid, exit)
+
+        # seed agent
+        feat_seed = self.seed_head(feat_a[-1:]) + self.seed_feature.weight[:, None]
+        next_state_prob_seed = self.seed_state_predict_head(feat_seed)
+        next_state_idx_seed = next_state_prob_seed.softmax(dim=-1).argmax(dim=-1, keepdim=True) # (self.seed_size, num_step, 1)
+
+        next_type_prob_seed = self.seed_type_predict_head(feat_seed)
+        next_type_idx_seed = next_type_prob_seed.softmax(dim=-1).argmax(dim=-1, keepdim=True)
+
+        next_type_index_gt = agent_type_index[:, None].expand(-1, num_step).roll(shifts=-1, dims=1)
+
+        # polygon token for bos token
+        # next_bos_pl_prob = self.bos_pl_predict_head(feat_a)
+        # next_bos_pl_prob_softmax = torch.softmax(next_bos_pl_prob, dim=-1)
+        # _, next_bos_pl_idx = torch.topk(next_bos_pl_prob_softmax, k=1, dim=-1) # (num_agent, num_step, 1)
+
+        # next_bos_pl_index_gt = agent_enter_pl_token_id.roll(shifts=-1, dims=-1)
+
+        # offset token for bos token
+        # next_bos_offset_prob = self.bos_offset_predict_head(feat_a)
+        # next_bos_offset_prob_softmax = torch.softmax(next_bos_offset_prob, dim=-1)
+        # _, next_bos_offset_idx = torch.topk(next_bos_offset_prob_softmax, k=1, dim=-1)
+
+        # next_bos_offset_index_gt = agent_enter_offset_token_idx.roll(shifts=-1, dims=-1)
+
+        # next token prediction mask
+        bos_token_index = torch.nonzero(agent_state_index == self.enter_state)
+        eos_token_index = torch.nonzero(agent_state_index == self.exit_state)
+
+        # mask for motion tokens
+        next_token_eval_mask = mask.clone()
+        next_token_eval_mask = next_token_eval_mask * next_token_eval_mask.roll(shifts=-1, dims=1) * next_token_eval_mask.roll(shifts=1, dims=1)
+        for bos_token_index_ in bos_token_index:
+            next_token_eval_mask[bos_token_index_[0], bos_token_index_[1] : bos_token_index_[1] + 1] = 1
+            next_token_eval_mask[bos_token_index_[0], bos_token_index_[1] + 1 : bos_token_index_[1] + 2] = \
+                            mask[bos_token_index_[0], bos_token_index_[1] + 2 : bos_token_index_[1] + 3]
+        next_token_eval_mask[eos_token_index[:, 0], eos_token_index[:, 1]] = 0
+
+        # mask for state tokens
+        next_state_eval_mask = mask.clone()
+        next_state_eval_mask = next_state_eval_mask * next_state_eval_mask.roll(shifts=-1, dims=1) * next_state_eval_mask.roll(shifts=1, dims=1)
+        for bos_token_index_ in bos_token_index:
+            next_state_eval_mask[bos_token_index_[0], :bos_token_index_[1]] = 0
+            next_state_eval_mask[bos_token_index_[0], bos_token_index_[1] : bos_token_index_[1] + 1] = 1
+            next_state_eval_mask[bos_token_index_[0], bos_token_index_[1] + 1 : bos_token_index_[1] + 2] = \
+                            mask[bos_token_index_[0], bos_token_index_[1] + 2 : bos_token_index_[1] + 3]
+        for eos_token_index_ in eos_token_index:
+            next_state_eval_mask[eos_token_index_[0], eos_token_index_[1] + 1:] = 1
+            next_state_eval_mask[eos_token_index_[0], eos_token_index_[1] : eos_token_index_[1] + 1] = \
+                            mask[eos_token_index_[0], eos_token_index_[1] - 1 : eos_token_index_[1]]
+
+        # seed agents
+        next_bos_token_index = torch.nonzero(next_state_index_gt == self.enter_state)
+        next_bos_token_index = next_bos_token_index[next_bos_token_index[:, 1] < num_step - 1]
+
+        next_state_index_gt_seed = torch.full((self.seed_size, num_step), self.seed_state_type.index('invalid'), device=next_state_index_gt.device)
+        next_type_index_gt_seed = torch.full((self.seed_size, num_step), self.seed_agent_type.index('seed'), device=next_state_index_gt.device)
+        next_eval_mask_seed = torch.ones_like(next_state_index_gt_seed)
+
+        num_seed = torch.zeros(num_step, device=next_state_index_gt.device).long()
+        for next_bos_token_index_ in next_bos_token_index:
+            if num_seed[next_bos_token_index_[1]] < self.seed_size:
+                next_state_index_gt_seed[num_seed[next_bos_token_index_[1]], next_bos_token_index_[1]] = self.seed_state_type.index('enter')
+                next_type_index_gt_seed[num_seed[next_bos_token_index_[1]], next_bos_token_index_[1]] = next_type_index_gt[next_bos_token_index_[0], next_bos_token_index_[1]]
+                num_seed[next_bos_token_index_[1]] += 1
+
+        # the last timestep is the beginning of the sequence (also the input)
+        next_token_eval_mask[:, -1] = 0
+        next_state_eval_mask[:, -1] = 0
+        next_eval_mask_seed[:, -1] = 0
+        # next_bos_token_eval_mask[:, -1] = False
+
+        # no invalid motion token will be supervised
+        if (next_token_index_gt[next_token_eval_mask] < 0).any():
+            raise RuntimeError()
+
+        next_state_index_gt[next_state_index_gt == self.exit_state] = self.valid_state_type.index('exit')
+
+        return {'x_a': feat_a,
+                # motion token
+                'next_token_idx': next_token_idx,
+                'next_token_prob': next_token_prob,
+                'next_token_idx_gt': next_token_index_gt,
+                'next_token_eval_mask': next_token_eval_mask.bool(),
+                # state token
+                'next_state_idx': next_state_idx,
+                'next_state_prob': next_state_prob,
+                'next_state_idx_gt': next_state_index_gt,
+                'next_state_eval_mask': next_state_eval_mask.bool(),
+                # seed agent
+                'next_state_idx_seed': next_state_idx_seed,
+                'next_state_prob_seed': next_state_prob_seed,
+                'next_state_idx_gt_seed': next_state_index_gt_seed,
+                'next_type_idx_seed': next_type_idx_seed,
+                'next_type_prob_seed': next_type_prob_seed,
+                'next_type_idx_gt_seed': next_type_index_gt_seed,
+                'next_eval_mask_seed': next_eval_mask_seed.bool(),
+                # pl token for bos
+                # 'next_bos_pl_idx': next_bos_pl_idx,
+                # 'next_bos_pl_prob': next_bos_pl_prob,
+                # 'next_bos_pl_index_gt': next_bos_pl_index_gt,
+                # offset token for bos
+                # 'next_bos_offset_idx': next_bos_offset_idx,
+                # 'next_bos_offset_prob': next_bos_offset_prob,
+                # 'next_bos_offset_index_gt': next_bos_offset_index_gt,
+                # 'next_bos_token_eval_mask': next_bos_token_eval_mask,
+                }
+
+    def inference(self,
+                  data: HeteroData,
+                  map_enc: Mapping[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
+
+        start_state_idx = data['agent']['state_idx'][:, (self.num_historical_steps - 1) // self.shift]
+        filter_mask = (start_state_idx == self.valid_state) | (start_state_idx == self.exit_state)
+        seed_step_mask = data['agent']['state_idx'][:, (self.num_historical_steps - 1) // self.shift:] == self.enter_state
+        seed_agent_index_per_step = [torch.nonzero(seed_step_mask[:, t]).squeeze(dim=-1) for t in range(seed_step_mask.shape[1])]
+        if torch.any(seed_step_mask.sum(dim=0) > self.seed_size):
+            raise RuntimeError(f"Seed size {self.seed_size} is too small.")
+
+        # num_historical_steps=11
+        eval_mask = data['agent']['valid_mask'][filter_mask, self.num_historical_steps - 1]
+
+        if self.predict_state:
+            eval_mask = torch.ones_like(eval_mask).bool()
+
+        # agent attributes
+        pos_a = data['agent']['token_pos'][filter_mask].clone() # (num_agent, num_step, 2)
+        state_a = data['agent']['state_idx'][filter_mask].clone() # (num_agent, num_step)
+        head_a = data['agent']['token_heading'][filter_mask].clone() # (num_agent, num_step)
+        gt_traj = data['agent']['position'][filter_mask, self.num_historical_steps:, :self.input_dim].contiguous()
+        num_agent, num_step, traj_dim = pos_a.shape
+
+        av_index = int(data['agent']['av_index'])
+        av_index -= (~filter_mask[:av_index]).sum()
+
+        # map attributes
+        pos_pl = data['pt_token']['position'][:, :2].clone() # (num_pl, 2)
+
+        # make future steps to zero
+        pos_a[:, (self.num_historical_steps - 1) // self.shift:] = 0
+        state_a[:, (self.num_historical_steps - 1) // self.shift:] = 0
+        head_a[:, (self.num_historical_steps - 1) // self.shift:] = 0
+
+        agent_valid_mask = data['agent']['raw_agent_valid_mask'][filter_mask].clone() # token_valid_mask
+        agent_valid_mask[:, (self.num_historical_steps - 1) // self.shift:] = True
+        agent_valid_mask[~eval_mask] = False
+        agent_token_index = data['agent']['token_idx'][filter_mask]
+        agent_state_index = data['agent']['state_idx'][filter_mask]
+        agent_type = data['agent']['type'][filter_mask]
+        agent_category = data['agent']['category'][filter_mask]
+
+        feat_a, head_vector_a, agent_token_traj_all, agent_token_emb, categorical_embs = self.agent_token_embedding(data,
+                                                                                                                    agent_token_index,
+                                                                                                                    agent_state_index,
+                                                                                                                    pos_a,
+                                                                                                                    head_a,
+                                                                                                                    inference=True,
+                                                                                                                    filter_mask=filter_mask,
+                                                                                                                    av_index=av_index,
+                                                                                            )
+        feat_seed = feat_a[-1:]
+        feat_a = feat_a[:-1]
+
+        agent_type = data["agent"]["type"][filter_mask]
+        veh_mask = agent_type == 0
+        cyc_mask = agent_type == 2
+        ped_mask = agent_type == 1
+
+        # self.num_recurrent_steps_val = 91 - 11 = 80
+        self.num_recurrent_steps_val = data["agent"]['position'].shape[1] - self.num_historical_steps
+        pred_traj = torch.zeros(pos_a.shape[0], self.num_recurrent_steps_val, 2, device=feat_a.device) # (num_agent, 80, 2)
+        pred_head = torch.zeros(pos_a.shape[0], self.num_recurrent_steps_val, device=feat_a.device)
+        pred_type = agent_type.clone()
+        pred_state = torch.zeros(pos_a.shape[0], self.num_recurrent_steps_val, device=feat_a.device)
+        pred_prob = torch.zeros(pos_a.shape[0], self.num_recurrent_steps_val // self.shift, device=feat_a.device) # (num_agent, 80 // 5 = 16)
+        next_token_idx_list = []
+        next_state_idx_list = []
+        next_bos_pl_idx_list = []
+        next_bos_offset_idx_list = []
+        feat_a_t_dict = {}
+        feat_sa_t_dict = {}
+
+        # build masks (init)
+        mask = agent_valid_mask.clone()
+        temporal_mask = mask.clone()
+        interact_mask = mask.clone()
+        if self.predict_state:
+
+            # find bos and eos index
+            is_bos = agent_state_index == self.enter_state
+            is_eos = agent_state_index == self.exit_state
+            bos_index = torch.where(is_bos.any(dim=1), torch.argmax(is_bos.long(), dim=1), torch.tensor(0))
+            eos_index = torch.where(is_eos.any(dim=1), torch.argmax(is_eos.long(), dim=1), torch.tensor(num_step - 1))
+
+            temporal_mask = torch.ones_like(mask)
+            motion_mask = torch.arange(mask.shape[1]).expand(mask.shape[0], mask.shape[1]).to(mask.device)
+            motion_mask = (motion_mask > bos_index[:, None]) & (motion_mask <= eos_index[:, None])
+            motion_mask[:, self.num_historical_steps // self.shift:] = False
+            temporal_mask[motion_mask] = mask[motion_mask]
+
+            interact_mask = torch.ones_like(mask)
+            non_motion_mask = ~motion_mask
+            non_motion_mask[:, self.num_historical_steps // self.shift:] = False
+            interact_mask[non_motion_mask] = False
+            interact_mask[agent_state_index == self.enter_state] = True
+
+        temporal_mask[:, (self.num_historical_steps - 1) // self.shift:] = True
+        interact_mask[:, (self.num_historical_steps - 1) // self.shift:] = True
+
+        # mapping network
+        # z = torch.randn(num_agent, self.hidden_dim).to(feat_a.device)
+        # w = self.mapping_network(z)
+
+        # we only need to predict 16 next tokens
+        for t in range(self.num_recurrent_steps_val // self.shift):
+
+            # feat_a = feat_a + w[:, None]
+            num_agent = pos_a.shape[0]
+
+            if t == 0:
+                inference_mask = temporal_mask.clone()
+                inference_mask = torch.cat([inference_mask, torch.ones_like(inference_mask[-1:])])
+                inference_mask[:, (self.num_historical_steps - 1) // self.shift + t:] = False
+            else:
+                inference_mask = torch.zeros_like(temporal_mask)
+                inference_mask = torch.cat([inference_mask, torch.zeros_like(inference_mask[-1:])])
+                inference_mask[:, max((self.num_historical_steps - 1) // self.shift + t - (self.num_interaction_steps // self.shift), 0) :
+                                      (self.num_historical_steps - 1) // self.shift + t] = True
+
+            interact_mask = torch.cat([interact_mask, torch.ones_like(interact_mask[:1])]).bool() # placeholder
+
+            edge_index_t, r_t = self.build_temporal_edge(pos_a, head_a, head_vector_a, state_a, temporal_mask, inference_mask,
+                                                         av_index=av_index)
+
+            # +1: placeholder for seed agent
+            batch_s = torch.arange(num_step, device=pos_a.device).repeat_interleave(num_agent + 1)
+            batch_pl = torch.arange(num_step, device=pos_a.device).repeat_interleave(data['pt_token']['num_nodes'])
+
+            # In the inference stage, we only infer the current stage for recurrent
+            edge_index_a2a, r_a2a = self.build_interaction_edge(pos_a, head_a, head_vector_a, state_a, batch_s,
+                                                                interact_mask, inference_mask, av_index=av_index)
+            edge_index_pl2a, r_pl2a = self.build_map2agent_edge(data, num_step, pos_a, head_a, head_vector_a, state_a, batch_s, batch_pl,
+                                                                interact_mask, inference_mask, av_index=av_index)
+            interact_mask = interact_mask[:-1]
+
+            # if t > 0:
+            #     feat_a_sum = feat_a.sum(dim=-1)
+            #     for a in range(pos_a.shape[0]):
+            #         t_1 = (self.num_historical_steps - 1) // self.shift + t - 1
+            #         print(f"agent {a} t_1 {t_1}")
+            #         print(f"token: {next_token_idx[a]}")
+            #         print(f"state: {next_state_idx[a]}")
+            #         print(f"feat_a_sum: {feat_a_sum[a, t_1]}")
+
+            for i in range(self.num_layers):
+
+                if (i in feat_a_t_dict) and (i in feat_sa_t_dict):
+                    feat_a = feat_a_t_dict[i]
+                    feat_seed = feat_sa_t_dict[i]
+
+                feat_a = torch.cat([feat_a, feat_seed], dim=0)
+
+                feat_a = feat_a.reshape(-1, self.hidden_dim)
+                feat_a = self.t_attn_layers[i](feat_a, r_t, edge_index_t)
+
+                feat_a = feat_a.reshape(-1, num_step, self.hidden_dim).transpose(0, 1).reshape(-1, self.hidden_dim)
+                feat_a = self.pt2a_attn_layers[i]((
+                    map_enc['x_pt'].repeat_interleave(repeats=num_step, dim=0).reshape(-1, num_step, self.hidden_dim).transpose(0, 1).reshape(
+                    -1, self.hidden_dim), feat_a), r_pl2a, edge_index_pl2a)
+
+                feat_a = self.a2a_attn_layers[i](feat_a, r_a2a, edge_index_a2a)
+                feat_a = feat_a.reshape(num_step, -1, self.hidden_dim).transpose(0, 1)
+
+                feat_seed = feat_a[-1:] # (1, num_step, hidden_dim)
+                feat_a = feat_a[:-1] # (num_agent, num_step, hidden_dim)
+
+                if t == 0:
+                    feat_a_t_dict[i + 1] = feat_a
+                    feat_sa_t_dict[i + 1] = feat_seed
+                else:
+                    # update agent features at current step
+                    n = feat_a_t_dict[i + 1].shape[0]
+                    feat_a_t_dict[i + 1][:n, (self.num_historical_steps - 1) // self.shift - 1 + t] = feat_a[:n, (self.num_historical_steps - 1) // self.shift - 1 + t]
+                    # add newly inserted agent features (only when t changed)
+                    if feat_a.shape[0] > n:
+                        m = feat_a.shape[0] - n
+                        feat_a_t_dict[i + 1] = torch.cat([feat_a_t_dict[i + 1], feat_a[-m:]])
+                    # update seed agent features at current step
+                    feat_sa_t_dict[i + 1][:, (self.num_historical_steps - 1) // self.shift - 1 + t] = feat_seed[:, (self.num_historical_steps - 1) // self.shift - 1 + t]
+
+            # next motion token
+            next_token_prob = self.token_predict_head(feat_a[:, (self.num_historical_steps - 1) // self.shift - 1 + t])
+            next_token_prob_softmax = torch.softmax(next_token_prob, dim=-1)
+            topk_token_prob, next_token_idx = torch.topk(next_token_prob_softmax, k=self.beam_size, dim=-1) # both (num_agent, beam_size) e.g. (31, 5)
+
+            # next state token
+            next_state_prob = self.state_predict_head(feat_a[:, (self.num_historical_steps - 1) // self.shift - 1 + t])
+            next_state_idx = next_state_prob.softmax(dim=-1).argmax(dim=-1)
+            next_state_idx[next_state_idx == self.valid_state_type.index('exit')] = self.exit_state
+
+            # seed agent
+            feat_seed = self.seed_head(feat_seed) + self.seed_feature.weight[:, None]
+            next_state_prob_seed = self.seed_state_predict_head(feat_seed[:, (self.num_historical_steps - 1) // self.shift - 1 + t])
+            next_state_idx_seed = next_state_prob_seed.softmax(dim=-1).argmax(dim=-1, keepdim=True)
+            next_state_idx_seed[next_state_idx_seed == self.seed_state_type.index('enter')] = self.enter_state
+
+            next_type_prob_seed = self.seed_type_predict_head(feat_seed[:, (self.num_historical_steps - 1) // self.shift - 1 + t])
+            next_type_idx_seed = next_type_prob_seed.softmax(dim=-1).argmax(dim=-1, keepdim=True)
+
+            # print(f"t: {t}")
+            # print(next_type_idx_seed[..., 0].tolist())
+
+            # bos pl prediction
+            # next_bos_pl_prob = self.bos_pl_predict_head(feat_a[:, (self.num_historical_steps - 1) // self.shift - 1 + t])
+            # next_bos_pl_prob_softmax = torch.softmax(next_bos_pl_prob, dim=-1)
+            # next_bos_pl_idx = torch.argmax(next_bos_pl_prob_softmax, dim=-1)
+
+            # bos offset prediction
+            # next_bos_offset_prob = self.bos_offset_predict_head(feat_a[:, (self.num_historical_steps - 1) // self.shift - 1 + t])
+            # next_bos_offset_prob_softmax = torch.softmax(next_bos_offset_prob, dim=-1)
+            # next_bos_offset_idx = torch.argmax(next_bos_offset_prob_softmax, dim=-1)
+
+            # convert the predicted token to a 0.5s (6 timesteps) trajectory
+            expanded_token_index = next_token_idx[..., None, None, None].expand(-1, -1, 6, 4, 2)
+            next_token_traj = torch.gather(agent_token_traj_all, 1, expanded_token_index) # (num_agent, beam_size, 6, 4, 2)
+
+            # apply rotation and translation on 'next_token_traj'
+            theta = head_a[:, (self.num_historical_steps - 1) // self.shift - 1 + t]
+            cos, sin = theta.cos(), theta.sin()
+            rot_mat = torch.zeros((num_agent, 2, 2), device=theta.device)
+            rot_mat[:, 0, 0] = cos
+            rot_mat[:, 0, 1] = sin
+            rot_mat[:, 1, 0] = -sin
+            rot_mat[:, 1, 1] = cos
+            agent_diff_rel = torch.bmm(next_token_traj.view(-1, 4, 2),
+                                       rot_mat[:, None, None, ...].repeat(1, self.beam_size, self.shift + 1, 1, 1).view(
+                                           -1, 2, 2)).view(num_agent, self.beam_size, self.shift + 1, 4, 2)
+            agent_pred_rel = agent_diff_rel + pos_a[:, (self.num_historical_steps - 1) // self.shift - 1 + t, :][:, None, None, None, ...]
+
+            # sample 1 most probable index of top beam_size tokens, (num_agent, beam_size) -> (num_agent, 1)
+            # then sample the agent_pred_rel, (num_agent, beam_size, 6, 4, 2) -> (num_agent, 6, 4, 2)
+            sample_token_index = torch.multinomial(topk_token_prob, 1).to(agent_pred_rel.device)
+            next_token_idx = next_token_idx.gather(dim=1, index=sample_token_index).squeeze(-1)
+            agent_pred_rel = agent_pred_rel.gather(dim=1,
+                                                   index=sample_token_index[..., None, None, None].expand(-1, -1, 6, 4,
+                                                                                                    2))[:, 0, ...]
+
+            # get predicted position and heading of current shifted timesteps
+            diff_xy = agent_pred_rel[:, 1:, 0, :] - agent_pred_rel[:, 1:, 3, :]
+            pred_traj[:num_agent, t * 5 : (t + 1) * 5] = agent_pred_rel[:, 1:, ...].clone().mean(dim=2)
+            pred_head[:num_agent, t * 5 : (t + 1) * 5] = torch.arctan2(diff_xy[:, :, 1], diff_xy[:, :, 0])
+            pred_state[:num_agent, t * 5 : (t + 1) * 5] = next_state_idx[:, None].repeat(1, 5)
+            # pred_prob[:num_agent, t] = topk_token_prob.gather(dim=-1, index=sample_token_index)[:, 0] # (num_agent, beam_size) -> (num_agent,)
+
+            # update pos/head/state of current step
+            pos_a[:, (self.num_historical_steps - 1) // self.shift + t] = agent_pred_rel[:, -1, ...].clone().mean(dim=1)
+            diff_xy = agent_pred_rel[:, -1, 0, :] - agent_pred_rel[:, -1, 3, :]
+            theta = torch.arctan2(diff_xy[:, 1], diff_xy[:, 0])
+            head_a[:, (self.num_historical_steps - 1) // self.shift + t] = theta
+            state_a[:, (self.num_historical_steps - 1) // self.shift + t] = next_state_idx
+
+            # the case that the current predicted state token is invalid/exit
+            is_eos = next_state_idx == self.exit_state
+            is_invalid = next_state_idx == self.invalid_state
+
+            next_token_idx[is_invalid] = -1
+            pos_a[is_invalid, (self.num_historical_steps - 1) // self.shift + t] = 0.
+            head_a[is_invalid, (self.num_historical_steps - 1) // self.shift + t] = 0.
+
+            mask[is_invalid, (self.num_historical_steps - 1) // self.shift + t] = False # to handle those newly-added agents
+            interact_mask[is_invalid, (self.num_historical_steps - 1) // self.shift + t] = False
+
+            agent_token_emb[is_invalid, (self.num_historical_steps - 1) // self.shift + t] = self.no_token_emb(torch.zeros(1, device=pos_a.device).long())
+
+            type_emb = categorical_embs[0].reshape(num_agent, num_step, -1)
+            shape_emb = categorical_embs[1].reshape(num_agent, num_step, -1)
+            type_emb[is_invalid, (self.num_historical_steps - 1) // self.shift + t] = self.type_a_emb(torch.tensor(self.all_agent_type.index('invalid'), device=pos_a.device).long())
+            shape_emb[is_invalid, (self.num_historical_steps - 1) // self.shift + t] = self.shape_emb(torch.full((1, 3), self.invalid_shape_value, device=pos_a.device))
+            categorical_embs = [type_emb.reshape(-1, self.hidden_dim), shape_emb.reshape(-1, self.hidden_dim)]
+
+            # FIXME: need to discuss!!!
+            # if is_eos.any():
+
+            #     pos_a[is_eos, (self.num_historical_steps - 1) // self.shift + t + 1:] = 0.
+            #     head_a[is_eos, (self.num_historical_steps - 1) // self.shift + t + 1:] = 0.
+            #     mask[is_eos, (self.num_historical_steps - 1) // self.shift + t + 1:] = False # to handle those newly-added agents
+            #     interact_mask[torch.cat([is_eos, torch.zeros(1, device=is_eos.device).bool()]), (self.num_historical_steps - 1) // self.shift + t + 1:] = False
+
+            #     agent_token_emb[is_eos, (self.num_historical_steps - 1) // self.shift + t + 1:] = self.no_token_emb(torch.zeros(1, device=pos_a.device).long())
+
+            #     type_emb = categorical_embs[0].reshape(num_agent, num_step, -1)
+            #     shape_emb = categorical_embs[1].reshape(num_agent, num_step, -1)
+            #     type_emb[is_eos, (self.num_historical_steps - 1) // self.shift + t + 1:] = self.type_a_emb(torch.tensor(self.all_agent_type.index('invalid'), device=pos_a.device).long())
+            #     shape_emb[is_eos, (self.num_historical_steps - 1) // self.shift + t + 1:] = self.shape_emb(torch.full((1, 3), self.invalid_shape_value, device=pos_a.device))
+            #     categorical_embs = [type_emb.reshape(-1, self.hidden_dim), shape_emb.reshape(-1, self.hidden_dim)]
+
+            # for sa in range(next_state_idx_seed.shape[0]):
+            #     if next_state_idx_seed[sa] == self.enter_state:
+            #         print(f"agent {sa} is entering at step {t}")
+
+            # insert new agents (from seed agent)
+            seed_agent_index_cur_step = seed_agent_index_per_step[t]
+            num_new_agent = min(len(seed_agent_index_cur_step), next_state_idx_seed.bool().sum())
+            new_agent_mask = next_state_idx_seed.bool()
+            next_state_idx_seed = next_state_idx_seed[new_agent_mask]
+            next_state_idx_seed = next_state_idx_seed[:num_new_agent]
+            next_type_idx_seed = next_type_idx_seed[new_agent_mask]
+            next_type_idx_seed = next_type_idx_seed[:num_new_agent]
+            selected_agent_index_cur_step = seed_agent_index_cur_step[:num_new_agent]
+            agent_token_index = torch.cat([agent_token_index, data['agent']['token_idx'][selected_agent_index_cur_step]])
+            agent_state_index = torch.cat([agent_state_index, data['agent']['state_idx'][selected_agent_index_cur_step]])
+            agent_category = torch.cat([agent_category, data['agent']['category'][selected_agent_index_cur_step]])
+            agent_valid_mask = torch.cat([agent_valid_mask, data['agent']['raw_agent_valid_mask'][selected_agent_index_cur_step]])
+            gt_traj = torch.cat([gt_traj, data['agent']['position'][selected_agent_index_cur_step, self.num_historical_steps:, :self.input_dim]])
+
+            # FIXME: under test!!! bos token index is -2
+            next_state_idx = torch.cat([next_state_idx, next_state_idx_seed], dim=0).long()
+            next_token_idx = torch.cat([next_token_idx, torch.zeros(num_new_agent, device=next_token_idx.device) - 2], dim=0).long()
+            mask = torch.cat([mask, torch.ones(num_new_agent, num_step, device=mask.device)], dim=0).bool()
+            temporal_mask = torch.cat([temporal_mask, torch.ones(num_new_agent, num_step, device=temporal_mask.device)], dim=0).bool()
+            interact_mask = torch.cat([interact_mask, torch.ones(num_new_agent, num_step, device=interact_mask.device)], dim=0).bool()
+
+            # new_pos_a = ego_pos_a[None].repeat(num_new_agent, 1, 1)
+            # new_head_a = ego_head_a[None].repeat(num_new_agent, 1)
+            new_pos_a = torch.zeros(num_new_agent, num_step, 2, device=pos_a.device)
+            new_head_a = torch.zeros(num_new_agent, num_step, device=pos_a.device)
+            new_state_a = torch.zeros(num_new_agent, num_step, device=state_a.device)
+            new_shape_a = torch.full((num_new_agent, num_step, 3), self.invalid_shape_value, device=pos_a.device)
+            new_type_a = torch.full((num_new_agent, num_step), self.all_agent_type.index('invalid'), device=pos_a.device)
+
+            if num_new_agent > 0:
+                gt_bos_pos_a = data['agent']['position'][seed_agent_index_cur_step[:num_new_agent], (self.num_historical_steps - 1) // self.shift + t]
+                new_pos_a[:, (self.num_historical_steps - 1) // self.shift + t] = gt_bos_pos_a[:, :2].clone()
+                pos_a = torch.cat([pos_a, new_pos_a], dim=0)
+
+                gt_bos_head_a = data['agent']['heading'][seed_agent_index_cur_step[:num_new_agent], (self.num_historical_steps - 1) // self.shift + t]
+                new_head_a[:, (self.num_historical_steps - 1) // self.shift + t] = gt_bos_head_a.clone()
+                head_a = torch.cat([head_a, new_head_a], dim=0)
+
+                gt_bos_shape_a = data['agent']['shape'][seed_agent_index_cur_step[:num_new_agent], self.num_historical_steps - 1]
+                gt_bos_type_a = data['agent']['type'][seed_agent_index_cur_step[:num_new_agent]]
+                new_shape_a[:, (self.num_historical_steps - 1) // self.shift + t:] = gt_bos_shape_a.clone()[:, None]
+                new_type_a[:, (self.num_historical_steps - 1) // self.shift + t:] = gt_bos_type_a.clone()[:, None]
+                # new_type_a[:, (self.num_historical_steps - 1) // self.shift + t] = next_type_idx_seed
+                pred_type = torch.cat([pred_type, new_type_a[:, (self.num_historical_steps - 1) // self.shift + t]])
+
+                new_state_a[:, (self.num_historical_steps - 1) // self.shift + t] = self.enter_state
+                state_a = torch.cat([state_a, new_state_a], dim=0)
+
+                mask[-num_new_agent:, :(self.num_historical_steps - 1) // self.shift + t + 1] = 0
+                interact_mask[-num_new_agent:, :(self.num_historical_steps - 1) // self.shift + t] = 0
+
+                # update all steps
+                new_pred_traj = torch.zeros(num_new_agent, self.num_recurrent_steps_val, 2, device=pos_a.device)
+                new_pred_traj[:, t * 5 : (t + 1) * 5] = new_pos_a[:, (self.num_historical_steps - 1) // self.shift + t][:, None].repeat(1, 5, 1)
+                pred_traj = torch.cat([pred_traj, new_pred_traj], dim=0)
+
+                new_pred_head = torch.zeros(num_new_agent, self.num_recurrent_steps_val, device=pos_a.device)
+                new_pred_head[:, t * 5 : (t + 1) * 5] = new_head_a[:, (self.num_historical_steps - 1) // self.shift + t][:, None].repeat(1, 5)
+                pred_head = torch.cat([pred_head, new_pred_head], dim=0)
+
+                new_pred_state = torch.zeros(num_new_agent, self.num_recurrent_steps_val, device=pos_a.device)
+                new_pred_state[:, t * 5 : (t + 1) * 5] = next_state_idx_seed[:, None].repeat(1, 5)
+                pred_state = torch.cat([pred_state, new_pred_state], dim=0)
+
+            # handle the position/heading of bos token
+            # bos_pl_pos = pos_pl[next_bos_pl_idx[is_bos].long()]
+            # bos_offset_pos = discretize_neighboring(neighbor_index=next_bos_offset_idx[is_bos])
+            # pos_a[is_bos, (self.num_historical_steps - 1) // self.shift + t] += (bos_pl_pos + bos_offset_pos)
+            # # headings before bos token remains 0 which align with training process
+            # head_a[is_bos, (self.num_historical_steps - 1) // self.shift + t] += 0.
+
+                # add new agents token embeddings
+                agent_token_emb = torch.cat([agent_token_emb, self.no_token_emb(torch.zeros(1, device=pos_a.device).long())[None, :].repeat(num_new_agent, num_step, 1)])
+                veh_mask = torch.cat([veh_mask, next_type_idx_seed == self.seed_agent_type.index('veh')])
+                ped_mask = torch.cat([ped_mask, next_type_idx_seed == self.seed_agent_type.index('ped')])
+                cyc_mask = torch.cat([cyc_mask, next_type_idx_seed == self.seed_agent_type.index('cyc')])
+
+                # add new agents trajectory embeddings
+                trajectory_token_veh = torch.from_numpy(self.trajectory_token['veh']).clone().to(pos_a.device).to(torch.float)
+                trajectory_token_ped = torch.from_numpy(self.trajectory_token['ped']).clone().to(pos_a.device).to(torch.float)
+                trajectory_token_cyc = torch.from_numpy(self.trajectory_token['cyc']).clone().to(pos_a.device).to(torch.float)
+
+                new_agent_token_traj_all = torch.zeros((num_new_agent, self.token_size, self.shift + 1, 4, 2), device=pos_a.device)
+                trajectory_token_all_veh = torch.from_numpy(self.trajectory_token_all['veh']).clone().to(pos_a.device).to(torch.float)
+                trajectory_token_all_ped = torch.from_numpy(self.trajectory_token_all['ped']).clone().to(pos_a.device).to(torch.float)
+                trajectory_token_all_cyc = torch.from_numpy(self.trajectory_token_all['cyc']).clone().to(pos_a.device).to(torch.float)
+                new_agent_token_traj_all[next_type_idx_seed == 0] = torch.cat(
+                    [trajectory_token_all_veh[:, :self.shift], trajectory_token_veh[:, None, ...]], dim=1)
+                new_agent_token_traj_all[next_type_idx_seed == 1] = torch.cat(
+                    [trajectory_token_all_ped[:, :self.shift], trajectory_token_ped[:, None, ...]], dim=1)
+                new_agent_token_traj_all[next_type_idx_seed == 2] = torch.cat(
+                    [trajectory_token_all_cyc[:, :self.shift], trajectory_token_cyc[:, None, ...]], dim=1)
+
+                agent_token_traj_all = torch.cat([agent_token_traj_all, new_agent_token_traj_all], dim=0)
+
+                # add new agents categorical embeddings
+                new_categorical_embs = [self.type_a_emb(new_type_a.reshape(-1).long()), self.shape_emb(new_shape_a.reshape(-1, 3))]
+                categorical_embs = [torch.cat([categorical_embs[0], new_categorical_embs[0]], dim=0),
+                                    torch.cat([categorical_embs[1], new_categorical_embs[1]], dim=0)]
+
+            # update token embeddings of current step
+            agent_token_emb[veh_mask, (self.num_historical_steps - 1) // self.shift + t] = self.agent_token_emb_veh[
+                next_token_idx[veh_mask]]
+            agent_token_emb[ped_mask, (self.num_historical_steps - 1) // self.shift + t] = self.agent_token_emb_ped[
+                next_token_idx[ped_mask]]
+            agent_token_emb[cyc_mask, (self.num_historical_steps - 1) // self.shift + t] = self.agent_token_emb_cyc[
+                next_token_idx[cyc_mask]]
+
+            motion_vector_a, head_vector_a = self.build_vector_a(pos_a, head_a, state_a)
+
+            motion_vector_a[:, (self.num_historical_steps - 1) // self.shift + 1 + t:] = 0.
+            head_vector_a[:, (self.num_historical_steps - 1) // self.shift + 1 + t:] = 0.
+            x_a = torch.stack(
+                [torch.norm(motion_vector_a[:, :, :2], p=2, dim=-1),
+                 angle_between_2d_vectors(ctr_vector=head_vector_a, nbr_vector=motion_vector_a[:, :, :2])], dim=-1)
+
+            x_b = x_a.clone()
+            x_a = self.x_a_emb(continuous_inputs=x_a.view(-1, x_a.size(-1)),
+                               categorical_embs=categorical_embs)
+            x_a = x_a.view(-1, num_step, self.hidden_dim)
+
+            s_a = self.state_a_emb(state_a.reshape(-1).long()).reshape(num_agent + num_new_agent, num_step, self.hidden_dim)
+            feat_a = torch.cat((agent_token_emb, x_a, s_a), dim=-1)
+            feat_a = self.fusion_emb(feat_a)
+
+            # if t >= 15:
+                # print(f"inference {t}")
+                # is_invalid = state_a == self.invalid_state
+                # is_bos = state_a == self.enter_state
+                # is_eos = state_a == self.exit_state
+                # bos_index = torch.where(is_bos.any(dim=1), torch.argmax(is_bos.long(), dim=1), torch.tensor(0))
+                # eos_index = torch.where(is_eos.any(dim=1), torch.argmax(is_eos.long(), dim=1), torch.tensor(num_step - 1))
+                # mask = torch.arange(num_step).expand(num_agent + num_new_agent, -1).to(state_a.device)
+                # mask = (mask >= bos_index[:, None]) & (mask <= eos_index[:, None] + 1)
+                # is_invalid[mask] = False
+                # is_invalid[:, (self.num_historical_steps - 1) // self.shift + 1 + t:] = False
+                # print(pos_a[:, :((self.num_historical_steps - 1) // self.shift + 1 + t)])
+                # print(state_a[:, :((self.num_historical_steps - 1) // self.shift + 1 + t)])
+                # print(pos_a[is_invalid][:, 0])
+                # print(head_a[is_invalid])
+                # print(categorical_embs[0].sum(dim=-1)[is_invalid.reshape(-1)])
+                # print(categorical_embs[1].sum(dim=-1)[is_invalid.reshape(-1)])
+                # print(motion_vector_a[is_invalid][:, 0])
+                # print(head_vector_a[is_invalid][:, 0])
+                # print(x_b.sum(dim=-1)[is_invalid])
+                # print(x_a.sum(dim=-1)[is_invalid])
+                # for a in range(state_a.shape[0]):
+                #     print(f"agent: {a}")
+                #     print(state_a[a])
+                #     print(is_invalid[a].long())
+                #     print(pos_a[a, :, 0])
+                #     print(motion_vector_a[a, :, 0])
+                # print(s_a.sum(dim=-1)[is_invalid])
+                # print(feat_a.sum(dim=-1)[is_invalid])
+
+            # replace the features of steps before bos of valid agents with the corresponding seed agent features
+            # is_bos = state_a == self.enter_state
+            # bos_index = torch.where(is_bos.any(dim=1), torch.argmax(is_bos.long(), dim=1), torch.tensor(num_step))
+            # before_bos_mask = torch.arange(num_step).expand(num_agent + num_new_agent, -1).to(state_a.device) < bos_index[:, None]
+            # feat_a[before_bos_mask] = feat_seed.repeat(num_agent + num_new_agent, 1, 1)[before_bos_mask]
+
+            # build seed agent features
+            motion_vector_seed = motion_vector_a[av_index : av_index + 1]
+            head_vector_seed = head_vector_a[av_index : av_index + 1]
+            feat_seed = self.build_invalid_agent_feature(num_step, pos_a.device, type_index=self.all_agent_type.index('seed'),
+                                                         motion_vector=motion_vector_seed, head_vector=head_vector_seed)
+            # print(f"inference {t}")
+            # print(feat_seed.sum(dim=-1))
+
+            next_token_idx_list.append(next_token_idx[:, None])
+            next_state_idx_list.append(next_state_idx[:, None])
+            # next_bos_pl_idx_list.append(next_bos_pl_idx[:, None])
+            # next_bos_offset_idx_list.append(next_bos_offset_idx[:, None])
+
+        # TODO: check this
+        # agent_valid_mask[agent_category != 3] = False
+
+        # print("inference")
+        # is_invalid = state_a == self.invalid_state
+        # is_bos = state_a == self.enter_state
+        # is_eos = state_a == self.exit_state
+        # bos_index = torch.where(is_bos.any(dim=1), torch.argmax(is_bos.long(), dim=1), torch.tensor(0))
+        # eos_index = torch.where(is_eos.any(dim=1), torch.argmax(is_eos.long(), dim=1), torch.tensor(num_step - 1))
+        # mask = torch.arange(num_step).expand(num_agent, -1).to(state_a.device)
+        # mask = (mask >= bos_index[:, None]) & (mask <= eos_index[:, None] + 1)
+        # is_invalid[mask] = False
+        # print(feat_a.sum(dim=-1)[is_invalid])
+        # print(pos_a[is_invalid][: 0])
+        # print(head_a[is_invalid])
+        # exit(1)
+
+        num_agent = pos_a.shape[0]
+        for i in range(len(next_token_idx_list)):
+            next_token_idx_list[i] = torch.cat([next_token_idx_list[i], torch.zeros(num_agent - next_token_idx_list[i].shape[0], 1, device=next_token_idx_list[i].device) - 1], dim=0).long()
+            next_state_idx_list[i] = torch.cat([next_state_idx_list[i], torch.zeros(num_agent - next_state_idx_list[i].shape[0], 1, device=next_state_idx_list[i].device)], dim=0).long()
+
+        # eval mask
+        next_token_eval_mask = agent_valid_mask.clone()
+        next_state_eval_mask = agent_valid_mask.clone()
+        bos_token_index = torch.nonzero(agent_state_index == self.enter_state)
+        eos_token_index = torch.nonzero(agent_state_index == self.exit_state)
+
+        next_token_eval_mask[bos_token_index[:, 0], bos_token_index[:, 1]] = 1
+
+        for bos_token_index_i in bos_token_index:
+            next_state_eval_mask[bos_token_index_i[0], :bos_token_index_i[1] + 2] = 1
+        for eos_token_index_i in eos_token_index:
+            next_state_eval_mask[eos_token_index_i[0], eos_token_index_i[1]:] = 1
+
+        # add history attributes
+        num_agent = pred_traj.shape[0]
+        num_init_agent = filter_mask.sum()
+
+        pred_traj = torch.cat([pred_traj, torch.zeros(num_agent, self.num_historical_steps - 1, *(pred_traj.shape[2:]), device=pred_traj.device)], dim=1)
+        pred_head = torch.cat([pred_head, torch.zeros(num_agent, self.num_historical_steps - 1, *(pred_head.shape[2:]), device=pred_head.device)], dim=1)
+        pred_state = torch.cat([pred_state, torch.zeros(num_agent, self.num_historical_steps - 1, *(pred_state.shape[2:]), device=pred_state.device)], dim=1)
+
+        pred_state[:num_init_agent, :self.num_historical_steps - 1] = data['agent']['state_idx'][filter_mask, :(self.num_historical_steps - 1) // self.shift].repeat_interleave(repeats=self.shift, dim=1)
+
+        historical_token_idx = data['agent']['token_idx'][filter_mask, :(self.num_historical_steps - 1) // self.shift]
+        historical_token_idx[historical_token_idx < 0] = 0
+        historical_token_traj_all = torch.gather(agent_token_traj_all, 1, historical_token_idx[..., None, None, None].expand(-1, -1, 6, 4, 2))
+        init_theta = head_a[:num_init_agent, 0]
+        cos, sin = init_theta.cos(), init_theta.sin()
+        rot_mat = torch.zeros((num_init_agent, 2, 2), device=init_theta.device)
+        rot_mat[:, 0, 0] = cos
+        rot_mat[:, 0, 1] = sin
+        rot_mat[:, 1, 0] = -sin
+        rot_mat[:, 1, 1] = cos
+        historical_token_traj_all = torch.bmm(historical_token_traj_all.view(-1, 4, 2),
+                                              rot_mat[:, None, None, ...].repeat(1, (self.num_historical_steps - 1) // self.shift, self.shift + 1, 1, 1).view(
+                                              -1, 2, 2)).view(num_init_agent, (self.num_historical_steps - 1) // self.shift, self.shift + 1, 4, 2)
+        historical_token_traj_all = historical_token_traj_all + pos_a[:num_init_agent, 0, :][:, None, None, None, ...]
+        pred_traj[:num_init_agent, :self.num_historical_steps - 1] = historical_token_traj_all[:, :, 1:, ...].clone().mean(dim=3).reshape(num_init_agent, -1, 2)
+        diff_xy = historical_token_traj_all[..., 1:, 0, :] - historical_token_traj_all[..., 1:, 3, :]
+        pred_head[:num_init_agent, :self.num_historical_steps - 1] = torch.arctan2(diff_xy[..., 1], diff_xy[..., 0]).reshape(num_init_agent, -1)
+
+        return {
+            'av_index': av_index,
+            'valid_mask': agent_valid_mask[:, self.num_historical_steps:],
+            'pos_a': pos_a[:, (self.num_historical_steps - 1) // self.shift:],
+            'head_a': head_a[:, (self.num_historical_steps - 1) // self.shift:],
+            'gt_traj': gt_traj,
+            'pred_traj': pred_traj,
+            'pred_head': pred_head,
+            'pred_type': list(map(lambda i: self.seed_agent_type[i], pred_type.tolist())),
+            'pred_state': pred_state,
+            'next_token_idx': torch.cat(next_token_idx_list, dim=-1), # (num_agent, num_step)
+            'next_token_idx_gt': agent_token_index,
+            'next_state_idx': torch.cat(next_state_idx_list, dim=-1) if len(next_state_idx_list) > 0 else None,
+            'next_state_idx_gt': agent_state_index,
+            'next_token_eval_mask': next_token_eval_mask,
+            'next_state_eval_mask': next_state_eval_mask,
+            # 'next_bos_pl_idx': torch.cat(next_bos_pl_idx_list, dim=-1),
+            # 'next_bos_offset_idx': torch.cat(next_bos_offset_idx_list, dim=-1),
+        }

seq_10_150_3_3_encode_occ_separate_offsets_bs8_128_no_seqindex_long_lastvalid_ablation_grid/backups/dev/modules/layers.py

@@ -0,0 +1,371 @@
+
+import math
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from typing import List, Optional, Tuple, Union
+from torch_geometric.nn.conv import MessagePassing
+from torch_geometric.utils import softmax
+
+from dev.utils.func import weight_init
+
+
+__all__ = ['AttentionLayer', 'FourierEmbedding', 'MLPEmbedding', 'MLPLayer', 'MappingNetwork']
+
+
+class AttentionLayer(MessagePassing):
+
+    def __init__(self,
+                 hidden_dim: int,
+                 num_heads: int,
+                 head_dim: int,
+                 dropout: float,
+                 bipartite: bool,
+                 has_pos_emb: bool,
+                 **kwargs) -> None:
+        super(AttentionLayer, self).__init__(aggr='add', node_dim=0, **kwargs)
+        self.num_heads = num_heads
+        self.head_dim = head_dim
+        self.has_pos_emb = has_pos_emb
+        self.scale = head_dim ** -0.5
+
+        self.to_q = nn.Linear(hidden_dim, head_dim * num_heads)
+        self.to_k = nn.Linear(hidden_dim, head_dim * num_heads, bias=False)
+        self.to_v = nn.Linear(hidden_dim, head_dim * num_heads)
+        if has_pos_emb:
+            self.to_k_r = nn.Linear(hidden_dim, head_dim * num_heads, bias=False)
+            self.to_v_r = nn.Linear(hidden_dim, head_dim * num_heads)
+        self.to_s = nn.Linear(hidden_dim, head_dim * num_heads)
+        self.to_g = nn.Linear(head_dim * num_heads + hidden_dim, head_dim * num_heads)
+        self.to_out = nn.Linear(head_dim * num_heads, hidden_dim)
+        self.attn_drop = nn.Dropout(dropout)
+        self.ff_mlp = nn.Sequential(
+            nn.Linear(hidden_dim, hidden_dim * 4),
+            nn.ReLU(inplace=True),
+            nn.Dropout(dropout),
+            nn.Linear(hidden_dim * 4, hidden_dim),
+        )
+        if bipartite:
+            self.attn_prenorm_x_src = nn.LayerNorm(hidden_dim)
+            self.attn_prenorm_x_dst = nn.LayerNorm(hidden_dim)
+        else:
+            self.attn_prenorm_x_src = nn.LayerNorm(hidden_dim)
+            self.attn_prenorm_x_dst = self.attn_prenorm_x_src
+        if has_pos_emb:
+            self.attn_prenorm_r = nn.LayerNorm(hidden_dim)
+        self.attn_postnorm = nn.LayerNorm(hidden_dim)
+        self.ff_prenorm = nn.LayerNorm(hidden_dim)
+        self.ff_postnorm = nn.LayerNorm(hidden_dim)
+        self.apply(weight_init)
+
+    def forward(self,
+                x: Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]],
+                r: Optional[torch.Tensor],
+                edge_index: torch.Tensor) -> torch.Tensor:
+        if isinstance(x, torch.Tensor):
+            x_src = x_dst = self.attn_prenorm_x_src(x)
+        else:
+            x_src, x_dst = x
+            x_src = self.attn_prenorm_x_src(x_src)
+            x_dst = self.attn_prenorm_x_dst(x_dst)
+            x = x[1]
+        if self.has_pos_emb and r is not None:
+            r = self.attn_prenorm_r(r)
+        x = x + self.attn_postnorm(self._attn_block(x_src, x_dst, r, edge_index))
+        x = x + self.ff_postnorm(self._ff_block(self.ff_prenorm(x)))
+        return x
+
+    def message(self,
+                q_i: torch.Tensor,
+                k_j: torch.Tensor,
+                v_j: torch.Tensor,
+                r: Optional[torch.Tensor],
+                index: torch.Tensor,
+                ptr: Optional[torch.Tensor]) -> torch.Tensor:
+        if self.has_pos_emb and r is not None:
+            k_j = k_j + self.to_k_r(r).view(-1, self.num_heads, self.head_dim)
+            v_j = v_j + self.to_v_r(r).view(-1, self.num_heads, self.head_dim)
+        sim = (q_i * k_j).sum(dim=-1) * self.scale
+        attn = softmax(sim, index, ptr)
+        self.attention_weight = attn.sum(-1).detach()
+        attn = self.attn_drop(attn)
+        return v_j * attn.unsqueeze(-1)
+
+    def update(self,
+               inputs: torch.Tensor,
+               x_dst: torch.Tensor) -> torch.Tensor:
+        inputs = inputs.view(-1, self.num_heads * self.head_dim)
+        g = torch.sigmoid(self.to_g(torch.cat([inputs, x_dst], dim=-1)))
+        return inputs + g * (self.to_s(x_dst) - inputs)
+
+    def _attn_block(self,
+                    x_src: torch.Tensor,
+                    x_dst: torch.Tensor,
+                    r: Optional[torch.Tensor],
+                    edge_index: torch.Tensor) -> torch.Tensor:
+        q = self.to_q(x_dst).view(-1, self.num_heads, self.head_dim)
+        k = self.to_k(x_src).view(-1, self.num_heads, self.head_dim)
+        v = self.to_v(x_src).view(-1, self.num_heads, self.head_dim)
+        agg = self.propagate(edge_index=edge_index, x_dst=x_dst, q=q, k=k, v=v, r=r)
+        return self.to_out(agg)
+
+    def _ff_block(self, x: torch.Tensor) -> torch.Tensor:
+        return self.ff_mlp(x)
+
+
+class FourierEmbedding(nn.Module):
+
+    def __init__(self,
+                 input_dim: int,
+                 hidden_dim: int,
+                 num_freq_bands: int) -> None:
+        super(FourierEmbedding, self).__init__()
+        self.input_dim = input_dim
+        self.hidden_dim = hidden_dim
+
+        self.freqs = nn.Embedding(input_dim, num_freq_bands) if input_dim != 0 else None
+        self.mlps = nn.ModuleList(
+            [nn.Sequential(
+                nn.Linear(num_freq_bands * 2 + 1, hidden_dim),
+                nn.LayerNorm(hidden_dim),
+                nn.ReLU(inplace=True),
+                nn.Linear(hidden_dim, hidden_dim),
+            )
+                for _ in range(input_dim)])
+        self.to_out = nn.Sequential(
+            nn.LayerNorm(hidden_dim),
+            nn.ReLU(inplace=True),
+            nn.Linear(hidden_dim, hidden_dim),
+        )
+        self.apply(weight_init)
+
+    def forward(self,
+                continuous_inputs: Optional[torch.Tensor] = None,
+                categorical_embs: Optional[List[torch.Tensor]] = None) -> torch.Tensor:
+        if continuous_inputs is None:
+            if categorical_embs is not None:
+                x = torch.stack(categorical_embs).sum(dim=0)
+            else:
+                raise ValueError('Both continuous_inputs and categorical_embs are None')
+        else:
+            x = continuous_inputs.unsqueeze(-1) * self.freqs.weight * 2 * math.pi
+            # Warning: if your data are noisy, don't use learnable sinusoidal embedding
+            x = torch.cat([x.cos(), x.sin(), continuous_inputs.unsqueeze(-1)], dim=-1)
+            continuous_embs: List[Optional[torch.Tensor]] = [None] * self.input_dim
+            for i in range(self.input_dim):
+                continuous_embs[i] = self.mlps[i](x[:, i])
+            x = torch.stack(continuous_embs).sum(dim=0)
+            if categorical_embs is not None:
+                x = x + torch.stack(categorical_embs).sum(dim=0)
+        return self.to_out(x)
+
+
+class MLPEmbedding(nn.Module):
+    def __init__(self,
+                 input_dim: int,
+                 hidden_dim: int) -> None:
+        super(MLPEmbedding, self).__init__()
+        self.input_dim = input_dim
+        self.hidden_dim = hidden_dim
+        self.mlp = nn.Sequential(
+            nn.Linear(input_dim, 128),
+            nn.LayerNorm(128),
+            nn.ReLU(inplace=True),
+            nn.Linear(128, hidden_dim),
+            nn.LayerNorm(hidden_dim),
+            nn.ReLU(inplace=True),
+            nn.Linear(hidden_dim, hidden_dim))
+        self.apply(weight_init)
+
+    def forward(self,
+                continuous_inputs: Optional[torch.Tensor] = None,
+                categorical_embs: Optional[List[torch.Tensor]] = None) -> torch.Tensor:
+        if continuous_inputs is None:
+            if categorical_embs is not None:
+                x = torch.stack(categorical_embs).sum(dim=0)
+            else:
+                raise ValueError('Both continuous_inputs and categorical_embs are None')
+        else:
+            x = self.mlp(continuous_inputs)
+            if categorical_embs is not None:
+                x = x + torch.stack(categorical_embs).sum(dim=0)
+        return x
+
+
+class MLPLayer(nn.Module):
+
+    def __init__(self,
+                 input_dim: int,
+                 hidden_dim: int=None,
+                 output_dim: int=None) -> None:
+        super(MLPLayer, self).__init__()
+
+        if hidden_dim is None:
+            hidden_dim = output_dim
+
+        self.mlp = nn.Sequential(
+            nn.Linear(input_dim, hidden_dim),
+            nn.LayerNorm(hidden_dim),
+            nn.ReLU(inplace=True),
+            nn.Linear(hidden_dim, output_dim),
+        )
+        self.apply(weight_init)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.mlp(x)
+
+
+class MappingNetwork(nn.Module):
+    def __init__(self, z_dim, w_dim, layer_dim=None, num_layers=8):
+        super().__init__()
+
+        if not layer_dim:
+            layer_dim = w_dim
+        layer_dims = [z_dim] + [layer_dim] * (num_layers - 1) + [w_dim]
+
+        layers = []
+        for i in range(num_layers):
+            layers.extend([
+                nn.Linear(layer_dims[i], layer_dims[i + 1]),
+                nn.LeakyReLU(),
+            ])
+        self.layers = nn.Sequential(*layers)
+    
+    def forward(self, z):
+        w = self.layers(z)
+        return w
+
+
+# class FocalLoss:
+#     def __init__(self, alpha: float=.25, gamma: float=2):
+#         self.alpha = alpha
+#         self.gamma = gamma
+
+#     def __call__(self, inputs, targets):
+#         prob = inputs.sigmoid()
+#         ce_loss = F.binary_cross_entropy_with_logits(inputs, targets.float(), reduction='none')
+#         p_t = prob * targets + (1 - prob) * (1 - targets)
+#         loss = ce_loss * ((1 - p_t) ** self.gamma)
+
+#         if self.alpha >= 0:
+#             alpha_t = self.alpha * targets + (1 - self.alpha) * (1 - targets)
+#             loss = alpha_t * loss
+
+#         return loss.mean()
+
+
+class FocalLoss(nn.Module):
+    """Focal Loss, as described in https://arxiv.org/abs/1708.02002.
+    It is essentially an enhancement to cross entropy loss and is
+    useful for classification tasks when there is a large class imbalance.
+    x is expected to contain raw, unnormalized scores for each class.
+    y is expected to contain class labels.
+    Shape:
+        - x: (batch_size, C) or (batch_size, C, d1, d2, ..., dK), K > 0.
+        - y: (batch_size,) or (batch_size, d1, d2, ..., dK), K > 0.
+    """
+
+    def __init__(
+        self,
+        alpha: Optional[torch.Tensor] = None,
+        gamma: float = 0.0,
+        reduction: str = "mean",
+        ignore_index: int = -100,
+    ):
+        """Constructor.
+        Args:
+            alpha (Tensor, optional): Weights for each class. Defaults to None.
+            gamma (float, optional): A constant, as described in the paper.
+                Defaults to 0.
+            reduction (str, optional): 'mean', 'sum' or 'none'.
+                Defaults to 'mean'.
+            ignore_index (int, optional): class label to ignore.
+                Defaults to -100.
+        """
+        if reduction not in ("mean", "sum", "none"):
+            raise ValueError('Reduction must be one of: "mean", "sum", "none".')
+
+        super().__init__()
+        self.alpha = alpha
+        self.gamma = gamma
+        self.ignore_index = ignore_index
+        self.reduction = reduction
+
+        self.nll_loss = nn.NLLLoss(
+            weight=alpha, reduction="none", ignore_index=ignore_index
+        )
+
+    def __repr__(self):
+        arg_keys = ["alpha", "gamma", "ignore_index", "reduction"]
+        arg_vals = [self.__dict__[k] for k in arg_keys]
+        arg_strs = [f"{k}={v}" for k, v in zip(arg_keys, arg_vals)]
+        arg_str = ", ".join(arg_strs)
+        return f"{type(self).__name__}({arg_str})"
+
+    def forward(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
+        if x.ndim > 2:
+            # (N, C, d1, d2, ..., dK) --> (N * d1 * ... * dK, C)
+            c = x.shape[1]
+            x = x.permute(0, *range(2, x.ndim), 1).reshape(-1, c)
+            # (N, d1, d2, ..., dK) --> (N * d1 * ... * dK,)
+            y = y.view(-1)
+
+        unignored_mask = y != self.ignore_index
+        y = y[unignored_mask]
+        if len(y) == 0:
+            return 0.0
+        x = x[unignored_mask]
+
+        # compute weighted cross entropy term: -alpha * log(pt)
+        # (alpha is already part of self.nll_loss)
+        log_p = F.log_softmax(x, dim=-1)
+        ce = self.nll_loss(log_p, y)
+
+        # get true class column from each row
+        all_rows = torch.arange(len(x))
+        log_pt = log_p[all_rows, y]
+
+        # compute focal term: (1 - pt)^gamma
+        pt = log_pt.exp()
+        focal_term = (1 - pt) ** self.gamma
+
+        # the full loss: -alpha * ((1 - pt)^gamma) * log(pt)
+        loss = focal_term * ce
+
+        if self.reduction == "mean":
+            loss = loss.mean()
+        elif self.reduction == "sum":
+            loss = loss.sum()
+
+        return loss
+
+
+class OccLoss(nn.Module):
+
+    # geo_scal_loss
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, pred, target, mask=None):
+
+        nonempty_probs = torch.sigmoid(pred)
+        empty_probs = 1 - nonempty_probs
+
+        if mask is None:
+            mask = torch.ones_like(target).bool()
+
+        nonempty_target = target == 1
+        nonempty_target = nonempty_target[mask].float()
+        nonempty_probs = nonempty_probs[mask]
+        empty_probs = empty_probs[mask]
+
+        intersection = (nonempty_target * nonempty_probs).sum()
+        precision = intersection / nonempty_probs.sum()
+        recall = intersection / nonempty_target.sum()
+        spec = ((1 - nonempty_target) * (empty_probs)).sum() / (1 - nonempty_target).sum()
+
+        return (
+            F.binary_cross_entropy(precision, torch.ones_like(precision))
+            + F.binary_cross_entropy(recall, torch.ones_like(recall))
+            + F.binary_cross_entropy(spec, torch.ones_like(spec))
+        )

seq_10_150_3_3_encode_occ_separate_offsets_bs8_128_no_seqindex_long_lastvalid_ablation_grid/backups/dev/modules/map_decoder.py

@@ -0,0 +1,130 @@
+from typing import Dict
+import torch
+import torch.nn as nn
+from torch_cluster import radius_graph
+from torch_geometric.data import Batch
+from torch_geometric.data import HeteroData
+from torch_geometric.utils import subgraph
+
+from dev.modules.layers import MLPLayer, AttentionLayer, FourierEmbedding, MLPEmbedding
+from dev.utils.func import weight_init, wrap_angle, angle_between_2d_vectors
+
+
+class SMARTMapDecoder(nn.Module):
+
+    def __init__(self,
+                 dataset: str,
+                 input_dim: int,
+                 hidden_dim: int,
+                 num_historical_steps: int,
+                 pl2pl_radius: float,
+                 num_freq_bands: int,
+                 num_layers: int,
+                 num_heads: int,
+                 head_dim: int,
+                 dropout: float,
+                 map_token) -> None:
+
+        super(SMARTMapDecoder, self).__init__()
+        self.dataset = dataset
+        self.input_dim = input_dim
+        self.hidden_dim = hidden_dim
+        self.num_historical_steps = num_historical_steps
+        self.pl2pl_radius = pl2pl_radius
+        self.num_freq_bands = num_freq_bands
+        self.num_layers = num_layers
+        self.num_heads = num_heads
+        self.head_dim = head_dim
+        self.dropout = dropout
+
+        if input_dim == 2:
+            input_dim_r_pt2pt = 3
+        elif input_dim == 3:
+            input_dim_r_pt2pt = 4
+        else:
+            raise ValueError('{} is not a valid dimension'.format(input_dim))
+
+        self.type_pt_emb = nn.Embedding(17, hidden_dim)
+        self.side_pt_emb = nn.Embedding(4, hidden_dim)
+        self.polygon_type_emb = nn.Embedding(4, hidden_dim)
+        self.light_pl_emb = nn.Embedding(4, hidden_dim)
+
+        self.r_pt2pt_emb = FourierEmbedding(input_dim=input_dim_r_pt2pt, hidden_dim=hidden_dim,
+                                            num_freq_bands=num_freq_bands)
+        self.pt2pt_layers = nn.ModuleList(
+            [AttentionLayer(hidden_dim=hidden_dim, num_heads=num_heads, head_dim=head_dim, dropout=dropout,
+                            bipartite=False, has_pos_emb=True) for _ in range(num_layers)]
+        )
+        self.token_size = 1024
+        self.token_predict_head = MLPLayer(input_dim=hidden_dim, hidden_dim=hidden_dim,
+                                           output_dim=self.token_size)
+        input_dim_token = 22
+        self.token_emb = MLPEmbedding(input_dim=input_dim_token, hidden_dim=hidden_dim)
+        self.map_token = map_token
+        self.apply(weight_init)
+        self.mask_pt = False
+
+    def maybe_autocast(self, dtype=torch.float32):
+        return torch.cuda.amp.autocast(dtype=dtype)
+
+    def forward(self, data: HeteroData) -> Dict[str, torch.Tensor]:
+        pt_valid_mask = data['pt_token']['pt_valid_mask']
+        pt_pred_mask = data['pt_token']['pt_pred_mask']
+        pt_target_mask = data['pt_token']['pt_target_mask']
+        mask_s = pt_valid_mask
+
+        pos_pt = data['pt_token']['position'][:, :self.input_dim].contiguous()
+        orient_pt = data['pt_token']['orientation'].contiguous()
+        orient_vector_pt = torch.stack([orient_pt.cos(), orient_pt.sin()], dim=-1)
+        token_sample_pt = self.map_token['traj_src'].to(pos_pt.device).to(torch.float)
+        pt_token_emb_src = self.token_emb(token_sample_pt.view(token_sample_pt.shape[0], -1))
+        pt_token_emb = pt_token_emb_src[data['pt_token']['token_idx']]
+
+        x_pt = pt_token_emb
+
+        token2pl = data[('pt_token', 'to', 'map_polygon')]['edge_index']
+        token_light_type = data['map_polygon']['light_type'][token2pl[1]]
+        x_pt_categorical_embs = [self.type_pt_emb(data['pt_token']['type'].long()),
+                                 self.polygon_type_emb(data['pt_token']['pl_type'].long()),
+                                 self.light_pl_emb(token_light_type.long()),]
+        x_pt = x_pt + torch.stack(x_pt_categorical_embs).sum(dim=0)
+        edge_index_pt2pt = radius_graph(x=pos_pt[:, :2], r=self.pl2pl_radius,
+                                        batch=data['pt_token']['batch'] if isinstance(data, Batch) else None,
+                                        loop=False, max_num_neighbors=100)
+        if self.mask_pt:
+            edge_index_pt2pt = subgraph(subset=mask_s, edge_index=edge_index_pt2pt)[0]
+        rel_pos_pt2pt = pos_pt[edge_index_pt2pt[0]] - pos_pt[edge_index_pt2pt[1]]
+        rel_orient_pt2pt = wrap_angle(orient_pt[edge_index_pt2pt[0]] - orient_pt[edge_index_pt2pt[1]])
+        if self.input_dim == 2:
+            r_pt2pt = torch.stack(
+                [torch.norm(rel_pos_pt2pt[:, :2], p=2, dim=-1),
+                 angle_between_2d_vectors(ctr_vector=orient_vector_pt[edge_index_pt2pt[1]],
+                                          nbr_vector=rel_pos_pt2pt[:, :2]),
+                 rel_orient_pt2pt], dim=-1)
+        elif self.input_dim == 3:
+            r_pt2pt = torch.stack(
+                [torch.norm(rel_pos_pt2pt[:, :2], p=2, dim=-1),
+                 angle_between_2d_vectors(ctr_vector=orient_vector_pt[edge_index_pt2pt[1]],
+                                          nbr_vector=rel_pos_pt2pt[:, :2]),
+                 rel_pos_pt2pt[:, -1],
+                 rel_orient_pt2pt], dim=-1)
+        else:
+            raise ValueError('{} is not a valid dimension'.format(self.input_dim))
+
+        # layers
+        r_pt2pt = self.r_pt2pt_emb(continuous_inputs=r_pt2pt, categorical_embs=None)
+        for i in range(self.num_layers):
+            x_pt = self.pt2pt_layers[i](x_pt, r_pt2pt, edge_index_pt2pt)
+
+        next_token_prob = self.token_predict_head(x_pt[pt_pred_mask])
+        next_token_prob_softmax = torch.softmax(next_token_prob, dim=-1)
+        _, next_token_idx = torch.topk(next_token_prob_softmax, k=10, dim=-1)
+        next_token_index_gt = data['pt_token']['token_idx'][pt_target_mask]
+
+        return {
+            'x_pt': x_pt,
+            'map_next_token_idx': next_token_idx,
+            'map_next_token_prob': next_token_prob,
+            'map_next_token_idx_gt': next_token_index_gt,
+            'map_next_token_eval_mask': pt_pred_mask[pt_pred_mask]
+        }

seq_10_150_3_3_encode_occ_separate_offsets_bs8_128_no_seqindex_long_lastvalid_ablation_grid/backups/dev/modules/occ_decoder.py

@@ -0,0 +1,927 @@
+import os
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from typing import Dict, Mapping, Optional, Literal
+from torch_cluster import radius, radius_graph
+from torch_geometric.data import HeteroData, Batch
+from torch_geometric.utils import dense_to_sparse, subgraph
+from scipy.optimize import linear_sum_assignment
+
+from dev.modules.attr_tokenizer import Attr_Tokenizer
+from dev.modules.layers import *
+from dev.utils.visualization import *
+from dev.utils.func import angle_between_2d_vectors, wrap_angle, weight_init
+
+
+class SMARTOccDecoder(nn.Module):
+
+    def __init__(self,
+                 dataset: str,
+                 input_dim: int,
+                 hidden_dim: int,
+                 num_historical_steps: int,
+                 time_span: Optional[int],
+                 pl2a_radius: float,
+                 pl2seed_radius: float,
+                 a2a_radius: float,
+                 a2sa_radius: float,
+                 pl2sa_radius: float,
+                 num_freq_bands: int,
+                 num_layers: int,
+                 num_heads: int,
+                 head_dim: int,
+                 dropout: float,
+                 token_data: Dict,
+                 token_size: int,
+                 special_token_index: list=[],
+                 attr_tokenizer: Attr_Tokenizer=None,
+                 predict_motion: bool=False,
+                 predict_state: bool=False,
+                 predict_map: bool=False,
+                 predict_occ: bool=False,
+                 state_token: Dict[str, int]=None,
+                 seed_size: int=5,
+                 buffer_size: int=32,
+                 loss_weight: dict=None,
+                 logger=None) -> None:
+
+        super(SMARTOccDecoder, self).__init__()
+        self.dataset = dataset
+        self.input_dim = input_dim
+        self.hidden_dim = hidden_dim
+        self.num_historical_steps = num_historical_steps
+        self.time_span = time_span if time_span is not None else num_historical_steps
+        self.pl2a_radius = pl2a_radius
+        self.pl2seed_radius = pl2seed_radius
+        self.a2a_radius = a2a_radius
+        self.a2sa_radius = a2sa_radius
+        self.pl2sa_radius = pl2sa_radius
+        self.num_freq_bands = num_freq_bands
+        self.num_layers = num_layers
+        self.num_heads = num_heads
+        self.head_dim = head_dim
+        self.dropout = dropout
+        self.special_token_index = special_token_index
+        self.predict_motion = predict_motion
+        self.predict_state = predict_state
+        self.predict_map = predict_map
+        self.predict_occ = predict_occ
+        self.loss_weight = loss_weight
+        self.logger = logger
+
+        self.attr_tokenizer = attr_tokenizer
+
+        # state tokens
+        self.state_type = list(state_token.keys())
+        self.state_token = state_token
+        self.invalid_state = int(state_token['invalid'])
+        self.valid_state = int(state_token['valid'])
+        self.enter_state = int(state_token['enter'])
+        self.exit_state = int(state_token['exit'])
+
+        self.seed_state_type = ['invalid', 'enter']
+        self.valid_state_type = ['invalid', 'valid', 'exit']
+
+        input_dim_r_pt2a = 3
+        input_dim_r_a2a = 3
+
+        self.seed_size = seed_size
+        self.buffer_size = buffer_size
+
+        self.agent_type = ['veh', 'ped', 'cyc', 'seed']
+        self.type_a_emb = nn.Embedding(len(self.agent_type), hidden_dim)
+        self.shape_emb = MLPEmbedding(input_dim=3, hidden_dim=hidden_dim)
+        self.state_a_emb = nn.Embedding(len(self.state_type), hidden_dim)
+        self.motion_gap = 1.
+        self.heading_gap = 1.
+        self.invalid_shape_value = .1
+        self.invalid_motion_value = -2.
+        self.invalid_head_value = -2.
+
+        self.r_pt2a_emb = FourierEmbedding(input_dim=input_dim_r_pt2a, hidden_dim=hidden_dim,
+                                           num_freq_bands=num_freq_bands)
+        self.r_a2a_emb = FourierEmbedding(input_dim=input_dim_r_a2a, hidden_dim=hidden_dim,
+                                          num_freq_bands=num_freq_bands)
+
+        self.token_size = token_size # 2048
+        self.grid_size = self.attr_tokenizer.grid_size
+        self.angle_size = self.attr_tokenizer.angle_size
+        self.agent_limit = 3
+        self.pt_limit = 10
+        self.grid_agent_occ_head = MLPLayer(input_dim=hidden_dim, hidden_dim=self.grid_size,
+                                            output_dim=self.agent_limit * self.grid_size)
+        self.grid_pt_occ_head = MLPLayer(input_dim=hidden_dim, hidden_dim=self.grid_size,
+                                         output_dim=self.pt_limit * self.grid_size)
+
+        # self.num_seed_feature = 1
+        # self.num_seed_feature = self.seed_size
+        self.num_seed_feature = 10
+
+        self.trajectory_token = token_data['token'] # dict('veh', 'ped', 'cyc') (2048, 4, 2)
+        self.trajectory_token_traj = token_data['traj'] # (2048, 6, 3)
+        self.trajectory_token_all = token_data['token_all'] # (2048, 6, 4, 2)
+        self.apply(weight_init)
+
+        self.shift = 5
+        self.beam_size = 5
+        self.hist_mask = True
+        self.temporal_attn_to_invalid = False
+        self.use_rel = False
+
+        # seed agent
+        self.temporal_attn_seed = False
+        self.seed_attn_to_av = True
+        self.seed_use_ego_motion = False
+
+    def transform_rel(self, token_traj, prev_pos, prev_heading=None):
+        if prev_heading is None:
+            diff_xy = prev_pos[:, :, -1, :] - prev_pos[:, :, -2, :]
+            prev_heading = torch.arctan2(diff_xy[:, :, 1], diff_xy[:, :, 0])
+
+        num_agent, num_step, traj_num, traj_dim = token_traj.shape
+        cos, sin = prev_heading.cos(), prev_heading.sin()
+        rot_mat = torch.zeros((num_agent, num_step, 2, 2), device=prev_heading.device)
+        rot_mat[:, :, 0, 0] = cos
+        rot_mat[:, :, 0, 1] = -sin
+        rot_mat[:, :, 1, 0] = sin
+        rot_mat[:, :, 1, 1] = cos
+        agent_diff_rel = torch.bmm(token_traj.view(-1, traj_num, 2), rot_mat.view(-1, 2, 2)).view(num_agent, num_step, traj_num, traj_dim)
+        agent_pred_rel = agent_diff_rel + prev_pos[:, :, -1:, :]
+        return agent_pred_rel
+
+    def _agent_token_embedding(self, data, agent_token_index, agent_state, agent_offset_token_idx, pos_a, head_a,
+                               inference=False, filter_mask=None, av_index=None):
+
+        if filter_mask is None:
+            filter_mask = torch.ones_like(agent_state[:, 2], dtype=torch.bool)
+
+        num_agent, num_step, traj_dim = pos_a.shape # traj_dim=2
+        agent_type = data['agent']['type'][filter_mask]
+        veh_mask = (agent_type == 0)
+        ped_mask = (agent_type == 1)
+        cyc_mask = (agent_type == 2)
+
+        motion_vector_a, head_vector_a = self._build_vector_a(pos_a, head_a, agent_state)
+
+        trajectory_token_veh = torch.from_numpy(self.trajectory_token['veh']).clone().to(pos_a.device).to(torch.float)
+        trajectory_token_ped = torch.from_numpy(self.trajectory_token['ped']).clone().to(pos_a.device).to(torch.float)
+        trajectory_token_cyc = torch.from_numpy(self.trajectory_token['cyc']).clone().to(pos_a.device).to(torch.float)
+        self.agent_token_emb_veh = self.token_emb_veh(trajectory_token_veh.view(trajectory_token_veh.shape[0], -1)) # (token_size, 8)
+        self.agent_token_emb_ped = self.token_emb_ped(trajectory_token_ped.view(trajectory_token_ped.shape[0], -1))
+        self.agent_token_emb_cyc = self.token_emb_cyc(trajectory_token_cyc.view(trajectory_token_cyc.shape[0], -1))
+
+        # add bos token embedding
+        self.agent_token_emb_veh = torch.cat([self.agent_token_emb_veh, self.bos_token_emb(torch.zeros(1, device=pos_a.device).long())])
+        self.agent_token_emb_ped = torch.cat([self.agent_token_emb_ped, self.bos_token_emb(torch.zeros(1, device=pos_a.device).long())])
+        self.agent_token_emb_cyc = torch.cat([self.agent_token_emb_cyc, self.bos_token_emb(torch.zeros(1, device=pos_a.device).long())])
+
+        # add invalid token embedding
+        self.agent_token_emb_veh = torch.cat([self.agent_token_emb_veh, self.no_token_emb(torch.zeros(1, device=pos_a.device).long())])
+        self.agent_token_emb_ped = torch.cat([self.agent_token_emb_ped, self.no_token_emb(torch.zeros(1, device=pos_a.device).long())])
+        self.agent_token_emb_cyc = torch.cat([self.agent_token_emb_cyc, self.no_token_emb(torch.zeros(1, device=pos_a.device).long())])
+
+        # self.grid_token_emb = self.token_emb_grid(torch.stack([self.attr_tokenizer.dist,
+        #                                                        self.attr_tokenizer.dir], dim=-1).to(pos_a.device))
+        self.grid_token_emb = self.token_emb_grid(self.attr_tokenizer.grid)
+        self.grid_token_emb = torch.cat([self.grid_token_emb, self.invalid_offset_token_emb(torch.zeros(1, device=pos_a.device).long())])
+
+        if inference:
+            agent_token_traj_all = torch.zeros((num_agent, self.token_size, self.shift + 1, 4, 2), device=pos_a.device)
+            trajectory_token_all_veh = torch.from_numpy(self.trajectory_token_all['veh']).clone().to(pos_a.device).to(torch.float)
+            trajectory_token_all_ped = torch.from_numpy(self.trajectory_token_all['ped']).clone().to(pos_a.device).to(torch.float)
+            trajectory_token_all_cyc = torch.from_numpy(self.trajectory_token_all['cyc']).clone().to(pos_a.device).to(torch.float)
+            agent_token_traj_all[veh_mask] = torch.cat(
+                [trajectory_token_all_veh[:, :self.shift], trajectory_token_veh[:, None, ...]], dim=1)
+            agent_token_traj_all[ped_mask] = torch.cat(
+                [trajectory_token_all_ped[:, :self.shift], trajectory_token_ped[:, None, ...]], dim=1)
+            agent_token_traj_all[cyc_mask] = torch.cat(
+                [trajectory_token_all_cyc[:, :self.shift], trajectory_token_cyc[:, None, ...]], dim=1)
+
+        # additional token embeddings are already added -> -1: invalid, -2: bos
+        agent_token_emb = torch.zeros((num_agent, num_step, self.hidden_dim), device=pos_a.device)
+        agent_token_emb[veh_mask] = self.agent_token_emb_veh[agent_token_index[veh_mask]]
+        agent_token_emb[ped_mask] = self.agent_token_emb_ped[agent_token_index[ped_mask]]
+        agent_token_emb[cyc_mask] = self.agent_token_emb_cyc[agent_token_index[cyc_mask]]
+
+        offset_token_emb = self.grid_token_emb[agent_offset_token_idx]
+
+        # 'vehicle', 'pedestrian', 'cyclist', 'background'
+        is_invalid = agent_state == self.invalid_state
+        agent_types = data['agent']['type'].clone()[filter_mask].long().repeat_interleave(repeats=num_step, dim=0)
+        agent_types[is_invalid.reshape(-1)] = self.agent_type.index('seed')
+        agent_shapes = data['agent']['shape'].clone()[filter_mask, self.num_historical_steps - 1, :].repeat_interleave(repeats=num_step, dim=0)
+        agent_shapes[is_invalid.reshape(-1)] = self.invalid_shape_value
+
+        # TODO: fix ego_pos in inference mode
+        offset_pos = pos_a - pos_a[av_index].repeat_interleave(repeats=data['batch_size_a'], dim=0)
+        feat_a, categorical_embs = self._build_agent_feature(num_step, pos_a.device,
+                                                             motion_vector_a,
+                                                             head_vector_a,
+                                                             agent_token_emb,
+                                                             offset_token_emb,
+                                                             offset_pos=offset_pos,
+                                                             type=agent_types,
+                                                             shape=agent_shapes,
+                                                             state=agent_state,
+                                                             n=num_agent)
+
+        if inference:
+            return feat_a, agent_token_traj_all, agent_token_emb, categorical_embs
+        else:
+            # seed agent feature
+            if self.seed_use_ego_motion:
+                motion_vector_seed = motion_vector_a[av_index].repeat_interleave(repeats=self.num_seed_feature, dim=0)
+                head_vector_seed = head_vector_a[av_index].repeat_interleave(repeats=self.num_seed_feature, dim=0)
+            else:
+                motion_vector_seed = head_vector_seed = None
+            feat_seed, _ = self._build_agent_feature(num_step, pos_a.device,
+                                                    motion_vector_seed,
+                                                    head_vector_seed,
+                                                    state_index=self.invalid_state,
+                                                    n=data.num_graphs * self.num_seed_feature)
+
+            feat_a = torch.cat([feat_a, feat_seed], dim=0) # (a + n, t, d)
+
+            return feat_a
+
+    def _build_vector_a(self, pos_a, head_a, state_a):
+        num_agent = pos_a.shape[0]
+
+        motion_vector_a = torch.cat([pos_a.new_zeros(num_agent, 1, self.input_dim),
+                                     pos_a[:, 1:] - pos_a[:, :-1]], dim=1)
+
+        motion_vector_a[state_a == self.invalid_state] = self.invalid_motion_value
+
+        # invalid -> valid
+        is_last_invalid = (state_a.roll(shifts=1, dims=1) == self.invalid_state) & (state_a != self.invalid_state)
+        is_last_invalid[:, 0] = state_a[:, 0] == self.enter_state
+        motion_vector_a[is_last_invalid] = self.motion_gap
+
+        # valid -> invalid
+        is_last_valid = (state_a.roll(shifts=1, dims=1) != self.invalid_state) & (state_a == self.invalid_state)
+        is_last_valid[:, 0] = False
+        motion_vector_a[is_last_valid] = -self.motion_gap
+
+        head_a[state_a == self.invalid_state] == self.invalid_head_value
+        head_vector_a = torch.stack([head_a.cos(), head_a.sin()], dim=-1)
+
+        return motion_vector_a, head_vector_a
+
+    def _build_agent_feature(self, num_step, device,
+                             motion_vector=None,
+                             head_vector=None,
+                             agent_token_emb=None,
+                             agent_grid_emb=None,
+                             offset_pos=None,
+                             type=None,
+                             shape=None,
+                             categorical_embs_a=None,
+                             state=None,
+                             state_index=None,
+                             n=1):
+
+        if agent_token_emb is None:
+            agent_token_emb = self.no_token_emb(torch.zeros(1, device=device).long())[:, None].repeat(n, num_step, 1)
+            if state is not None:
+                agent_token_emb[state == self.enter_state] = self.bos_token_emb(torch.zeros(1, device=device).long())
+
+        if agent_grid_emb is None:
+            agent_grid_emb = self.grid_token_emb[None, None, self.grid_size // 2].repeat(n, num_step, 1)
+
+        if motion_vector is None or head_vector is None:
+            pos_a = torch.zeros((n, num_step, 2), device=device)
+            head_a = torch.zeros((n, num_step), device=device)
+            if state is None:
+                state = torch.full((n, num_step), self.invalid_state, device=device)
+            motion_vector, head_vector = self._build_vector_a(pos_a, head_a, state)
+
+        if offset_pos is None:
+            offset_pos = torch.zeros_like(motion_vector)
+
+        feature_a = torch.stack(
+            [torch.norm(motion_vector[:, :, :2], p=2, dim=-1),
+             angle_between_2d_vectors(ctr_vector=head_vector, nbr_vector=motion_vector[:, :, :2]),
+             # torch.norm(offset_pos[:, :, :2], p=2, dim=-1),
+             ], dim=-1)
+
+        if categorical_embs_a is None:
+            if type is None:
+                type = torch.tensor([self.agent_type.index('seed')], device=device)
+            if shape is None:
+                shape = torch.full((1, 3), self.invalid_shape_value, device=device)
+
+            categorical_embs_a = [self.type_a_emb(type.reshape(-1)), self.shape_emb(shape.reshape(-1, shape.shape[-1]))]
+
+        x_a = self.x_a_emb(continuous_inputs=feature_a.view(-1, feature_a.size(-1)),
+                           categorical_embs=categorical_embs_a)
+        x_a = x_a.view(-1, num_step, self.hidden_dim) # (a, t, d)
+
+        if state is None:
+            assert state_index is not None, f"state index need to be set when state tensor is None!"
+            state = torch.tensor([state_index], device=device)[:, None].repeat(n, num_step, 1) # do not use `expand`
+        s_a = self.state_a_emb(state.reshape(-1).long()).reshape(n, num_step, self.hidden_dim)
+
+        feat_a = torch.cat((agent_token_emb, x_a, s_a, agent_grid_emb), dim=-1)
+        feat_a = self.fusion_emb(feat_a) # (a, t, d)
+
+        return feat_a, categorical_embs_a
+
+    def _pad_feat(self, num_graph, av_index, *feats, num_seed_feature=None):
+
+        if num_seed_feature is None:
+            num_seed_feature = self.num_seed_feature
+
+        padded_feats = tuple()
+        for i in range(len(feats)):
+            padded_feats += (torch.cat([feats[i], feats[i][av_index].repeat_interleave(
+                repeats=num_seed_feature, dim=0)],
+                dim=0
+            ),)
+
+        pad_mask = torch.ones(*padded_feats[0].shape[:2], device=feats[0].device).bool() # (a, t)
+        pad_mask[-num_graph * num_seed_feature:] = False
+
+        return padded_feats + (pad_mask,)
+
+    def _build_seed_feat(self, data, pos_a, head_a, state_a, head_vector_a, mask, sort_indices, av_index):
+        seed_mask = sort_indices != av_index.repeat_interleave(repeats=data['batch_size_a'], dim=0)[:, None]
+        # TODO: fix batch_size!!!
+        print(mask.shape, sort_indices.shape, seed_mask.shape)
+        mask[-data.num_graphs * self.num_seed_feature:] = seed_mask[:self.num_seed_feature]
+
+        insert_pos_a = torch.gather(pos_a, dim=0, index=sort_indices[:self.num_seed_feature, :, None].expand(-1, -1, pos_a.shape[-1]))
+        pos_a[mask] = insert_pos_a[mask[-self.num_seed_feature:]]
+
+        state_a[-data.num_graphs * self.num_seed_feature:] = self.enter_state
+
+        return pos_a, head_a, state_a, head_vector_a, mask
+
+    def _build_temporal_edge(self, data, pos_a, head_a, state_a, head_vector_a, mask, inference_mask=None):
+
+        num_graph = data.num_graphs
+        num_agent = pos_a.shape[0]
+        hist_mask = mask.clone()
+
+        if not self.temporal_attn_to_invalid:
+            is_bos = state_a == self.enter_state
+            bos_index = torch.where(is_bos.any(dim=1), torch.argmax(is_bos.long(), dim=1), torch.tensor(0))
+            history_invalid_mask = torch.arange(mask.shape[1]).expand(mask.shape[0], mask.shape[1]).to(mask.device)
+            history_invalid_mask = (history_invalid_mask < bos_index[:, None])
+            hist_mask[history_invalid_mask] = False
+
+        if not self.temporal_attn_seed:
+            hist_mask[-num_graph * self.num_seed_feature:] = False
+            if inference_mask is not None:
+                inference_mask[-num_graph * self.num_seed_feature:] = False
+        else:
+            # WARNING: if use temporal attn to seed
+            # we need to fix the pos/head of seed!!!
+            raise RuntimeError("Wrong settings!")
+
+        pos_t = pos_a.reshape(-1, self.input_dim) # (num_agent * num_step, ...)
+        head_t = head_a.reshape(-1)
+        head_vector_t = head_vector_a.reshape(-1, 2)
+
+        # for those invalid agents won't predict any motion token, we don't attend to them
+        is_bos = state_a == self.enter_state
+        is_bos[-num_graph * self.num_seed_feature:] = False
+        bos_index = torch.where(is_bos.any(dim=1), torch.argmax(is_bos.long(), dim=1), torch.tensor(0))
+        motion_predict_start_index = torch.clamp(bos_index - self.time_span / self.shift + 1, min=0)
+        motion_predict_mask = torch.arange(hist_mask.shape[1]).expand(hist_mask.shape[0], -1).to(hist_mask.device)
+        motion_predict_mask = motion_predict_mask >= motion_predict_start_index[:, None]
+        hist_mask[~motion_predict_mask] = False
+
+        if self.hist_mask and self.training:
+            hist_mask[
+                torch.arange(mask.shape[0]).unsqueeze(1), torch.randint(0, mask.shape[1], (num_agent, 10))] = False
+            mask_t = hist_mask.unsqueeze(2) & hist_mask.unsqueeze(1)
+        elif inference_mask is not None:
+            mask_t = hist_mask.unsqueeze(2) & inference_mask.unsqueeze(1)
+        else:
+            mask_t = hist_mask.unsqueeze(2) & hist_mask.unsqueeze(1)
+
+        # mask_t: (num_agent, 18, 18), edge_index_t: (2, num_edge)
+        edge_index_t = dense_to_sparse(mask_t)[0]
+        edge_index_t = edge_index_t[:, (edge_index_t[1] - edge_index_t[0] > 0) &
+                                       (edge_index_t[1] - edge_index_t[0] <= self.time_span / self.shift)]
+        rel_pos_t = pos_t[edge_index_t[0]] - pos_t[edge_index_t[1]]
+        rel_head_t = wrap_angle(head_t[edge_index_t[0]] - head_t[edge_index_t[1]])
+
+        # handle the invalid steps
+        is_invalid = state_a == self.invalid_state
+        is_invalid_t = is_invalid.reshape(-1)
+
+        rel_pos_t[is_invalid_t[edge_index_t[0]] & ~is_invalid_t[edge_index_t[1]]] = -self.motion_gap
+        rel_pos_t[~is_invalid_t[edge_index_t[0]] & is_invalid_t[edge_index_t[1]]] = self.motion_gap
+        rel_head_t[is_invalid_t[edge_index_t[0]] & ~is_invalid_t[edge_index_t[1]]] = -self.heading_gap
+        rel_head_t[~is_invalid_t[edge_index_t[1]] & is_invalid_t[edge_index_t[1]]] = self.heading_gap
+
+        rel_pos_t[is_invalid_t[edge_index_t[0]] & is_invalid_t[edge_index_t[1]]] = self.invalid_motion_value
+        rel_head_t[is_invalid_t[edge_index_t[0]] & is_invalid_t[edge_index_t[1]]] = self.invalid_head_value
+
+        r_t = torch.stack(
+            [torch.norm(rel_pos_t[:, :2], p=2, dim=-1),
+             angle_between_2d_vectors(ctr_vector=head_vector_t[edge_index_t[1]], nbr_vector=rel_pos_t[:, :2]),
+             rel_head_t,
+             edge_index_t[0] - edge_index_t[1]], dim=-1)
+        r_t = self.r_t_emb(continuous_inputs=r_t, categorical_embs=None)
+
+        return edge_index_t, r_t
+
+    def _build_interaction_edge(self, data, pos_a, head_a, state_a, head_vector_a, batch_s, mask, pad_mask=None, inference_mask=None,
+                                av_index=None, seq_mask=None, seq_index=None, grid_index_a=None, **plot_kwargs):
+        num_graph = data.num_graphs
+        num_agent, num_step, _ = pos_a.shape
+        is_training = inference_mask is None
+
+        mask_a = mask.clone()
+
+        if pad_mask is None:
+            pad_mask = torch.ones_like(state_a).bool()
+
+        pos_s = pos_a.transpose(0, 1).reshape(-1, self.input_dim)
+        head_s = head_a.transpose(0, 1).reshape(-1)
+        head_vector_s = head_vector_a.transpose(0, 1).reshape(-1, 2)
+        pad_mask_s = pad_mask.transpose(0, 1).reshape(-1)
+        if inference_mask is not None:
+            mask_a = mask_a & inference_mask
+        mask_s = mask_a.transpose(0, 1).reshape(-1)
+
+        # build agent2agent bilateral connection
+        edge_index_a2a = radius_graph(x=pos_s[:, :2], r=self.a2a_radius, batch=batch_s, loop=False,
+                                      max_num_neighbors=300)
+        edge_index_a2a = subgraph(subset=mask_s & pad_mask_s, edge_index=edge_index_a2a)[0]
+
+        if os.getenv('PLOT_EDGE', False):
+            plot_interact_edge(edge_index_a2a, data['scenario_id'], data['batch_size_a'].cpu(), self.num_seed_feature, num_step,
+                               av_index=av_index, **plot_kwargs)
+
+        rel_pos_a2a = pos_s[edge_index_a2a[0]] - pos_s[edge_index_a2a[1]]
+        rel_head_a2a = wrap_angle(head_s[edge_index_a2a[0]] - head_s[edge_index_a2a[1]])
+
+        # handle the invalid steps
+        is_invalid = state_a == self.invalid_state
+        is_invalid_s = is_invalid.transpose(0, 1).reshape(-1)
+
+        rel_pos_a2a[is_invalid_s[edge_index_a2a[0]] & ~is_invalid_s[edge_index_a2a[1]]] = -self.motion_gap
+        rel_pos_a2a[~is_invalid_s[edge_index_a2a[0]] & is_invalid_s[edge_index_a2a[1]]] = self.motion_gap
+        rel_head_a2a[is_invalid_s[edge_index_a2a[0]] & ~is_invalid_s[edge_index_a2a[1]]] = -self.heading_gap
+        rel_head_a2a[~is_invalid_s[edge_index_a2a[1]] & is_invalid_s[edge_index_a2a[1]]] = self.heading_gap
+
+        rel_pos_a2a[is_invalid_s[edge_index_a2a[0]] & is_invalid_s[edge_index_a2a[1]]] = self.invalid_motion_value
+        rel_head_a2a[is_invalid_s[edge_index_a2a[0]] & is_invalid_s[edge_index_a2a[1]]] = self.invalid_head_value
+
+        r_a2a = torch.stack(
+            [torch.norm(rel_pos_a2a[:, :2], p=2, dim=-1),
+             angle_between_2d_vectors(ctr_vector=head_vector_s[edge_index_a2a[1]], nbr_vector=rel_pos_a2a[:, :2]),
+             rel_head_a2a,
+             torch.zeros_like(edge_index_a2a[0])], dim=-1)
+        r_a2a = self.r_a2a_emb(continuous_inputs=r_a2a, categorical_embs=None)
+
+        # add the edges which connect seed agents
+        if is_training:
+            mask_av = torch.ones_like(mask_a).bool()
+            if not self.seed_attn_to_av:
+                mask_av[av_index] = False
+            mask_a &= mask_av
+            edge_index_seed2a, r_seed2a = self._build_a2sa_edge(data, pos_a, head_a, head_vector_a, batch_s,
+                                                                mask_a.clone(), ~pad_mask.clone(), inference_mask=inference_mask,
+                                                                r=self.pl2seed_radius, max_num_neighbors=300,
+                                                                seq_mask=seq_mask, seq_index=seq_index, grid_index_a=grid_index_a, mode='grid')
+
+            if os.getenv('PLOT_EDGE', False):
+                plot_interact_edge(edge_index_seed2a, data['scenario_id'], data['batch_size_a'].cpu(), self.num_seed_feature, num_step,
+                                   'interact_edge_map_seed', av_index=av_index, **plot_kwargs)
+
+            edge_index_a2a = torch.cat([edge_index_a2a, edge_index_seed2a], dim=-1)
+            r_a2a = torch.cat([r_a2a, r_seed2a])
+
+            return edge_index_a2a, r_a2a, (edge_index_a2a.shape[1], edge_index_seed2a.shape[1]) #, nearest_dict
+
+        return edge_index_a2a, r_a2a
+
+    def _build_map2agent_edge(self, data, pos_a, head_a, state_a, head_vector_a, batch_s, batch_pl,
+                              mask, pad_mask=None, inference_mask=None, av_index=None, **kwargs):
+        num_graph = data.num_graphs
+        num_agent, num_step, _ = pos_a.shape
+        is_training = inference_mask is None
+
+        mask_pl2a = mask.clone()
+
+        if pad_mask is None:
+            pad_mask = torch.ones_like(state_a).bool()
+
+        pos_s = pos_a.transpose(0, 1).reshape(-1, self.input_dim)
+        head_s = head_a.transpose(0, 1).reshape(-1)
+        head_vector_s = head_vector_a.transpose(0, 1).reshape(-1, 2)
+        pad_mask_s = pad_mask.transpose(0, 1).reshape(-1)
+        if inference_mask is not None:
+            mask_pl2a = mask_pl2a & inference_mask
+        mask_pl2a = mask_pl2a.transpose(0, 1).reshape(-1)
+
+        ori_pos_pl = data['pt_token']['position'][:, :self.input_dim].contiguous()
+        ori_orient_pl = data['pt_token']['orientation'].contiguous()
+        pos_pl = ori_pos_pl.repeat(num_step, 1) # not `repeat_interleave`
+        orient_pl = ori_orient_pl.repeat(num_step)
+
+        # build map2agent directed graph
+        # edge_index_pl2a = radius(x=pos_s[:, :2], y=pos_pl[:, :2], r=self.pl2a_radius,
+        #                          batch_x=batch_s, batch_y=batch_pl, max_num_neighbors=300)
+        edge_index_pl2a = radius(x=pos_pl[:, :2], y=pos_s[:, :2], r=self.pl2a_radius,
+                                 batch_x=batch_pl, batch_y=batch_s, max_num_neighbors=5)
+        edge_index_pl2a = edge_index_pl2a[[1, 0]]
+        edge_index_pl2a = edge_index_pl2a[:, mask_pl2a[edge_index_pl2a[1]] &
+                                             pad_mask_s[edge_index_pl2a[1]]]
+
+        rel_pos_pl2a = pos_pl[edge_index_pl2a[0]] - pos_s[edge_index_pl2a[1]]
+        rel_orient_pl2a = wrap_angle(orient_pl[edge_index_pl2a[0]] - head_s[edge_index_pl2a[1]])
+
+        # handle the invalid steps
+        is_invalid = state_a == self.invalid_state
+        is_invalid_s = is_invalid.transpose(0, 1).reshape(-1)
+        rel_pos_pl2a[is_invalid_s[edge_index_pl2a[1]]] = self.motion_gap
+        rel_orient_pl2a[is_invalid_s[edge_index_pl2a[1]]] = self.heading_gap
+
+        r_pl2a = torch.stack(
+            [torch.norm(rel_pos_pl2a[:, :2], p=2, dim=-1),
+             angle_between_2d_vectors(ctr_vector=head_vector_s[edge_index_pl2a[1]], nbr_vector=rel_pos_pl2a[:, :2]),
+             rel_orient_pl2a], dim=-1)
+        r_pl2a = self.r_pt2a_emb(continuous_inputs=r_pl2a, categorical_embs=None)
+
+        # add the edges which connect seed agents
+        if is_training:
+            edge_index_pl2seed, r_pl2seed = self._build_map2sa_edge(data, pos_a, head_a, head_vector_a, batch_s, batch_pl,
+                                                                    ~pad_mask.clone(), inference_mask=inference_mask,
+                                                                    r=self.pl2seed_radius, max_num_neighbors=2048, mode='grid')
+
+            # sanity check
+            # pl2a_index = torch.zeros(pos_a.shape[0], num_step)
+            # pl2a_r = torch.zeros(pos_a.shape[0], num_step)
+            # for src_index in torch.unique(edge_index_pl2seed[1]):
+            #     src_row = src_index % pos_a.shape[0]
+            #     src_col = src_index // pos_a.shape[0]
+            #     pl2a_index[src_row, src_col] = edge_index_pl2seed[0, edge_index_pl2seed[1] == src_index].sum()
+            #     pl2a_r[src_row, src_col] = r_pl2seed[edge_index_pl2seed[1] == src_index].sum()
+            # print(pl2a_index)
+            # print(pl2a_r)
+            # exit(1)
+
+            if os.getenv('PLOT_EDGE', False):
+                plot_interact_edge(edge_index_pl2seed, data['scenario_id'], data['batch_size_a'].cpu(), self.num_seed_feature, num_step,
+                                'interact_edge_map_seed', av_index=av_index)
+
+            edge_index_pl2a = torch.cat([edge_index_pl2a, edge_index_pl2seed], dim=-1)
+            r_pl2a = torch.cat([r_pl2a, r_pl2seed])
+
+            return edge_index_pl2a, r_pl2a, (edge_index_pl2a.shape[1], edge_index_pl2seed.shape[1])
+
+        return edge_index_pl2a, r_pl2a
+
+    def _build_a2sa_edge(self, data, pos_a, head_a, head_vector_a, batch_s, mask_a, mask_sa,
+                         inference_mask=None, r=None, max_num_neighbors=8, seq_mask=None, seq_index=None,
+                         grid_index_a=None, mode: Literal['grid', 'heading']='heading', **plot_kwargs):
+
+        num_agent, num_step, _ = pos_a.shape
+        is_training = inference_mask is None
+
+        pos_s = pos_a.transpose(0, 1).reshape(-1, self.input_dim)
+        head_s = head_a.transpose(0, 1).reshape(-1)
+        head_vector_s = head_vector_a.transpose(0, 1).reshape(-1, 2)
+        if inference_mask is not None:
+            mask_a = mask_a & inference_mask
+            mask_sa = mask_sa & inference_mask
+        mask_s = mask_a.transpose(0, 1).reshape(-1)
+        mask_s_sa = mask_sa.transpose(0, 1).reshape(-1)
+
+        # build seed_agent2agent unilateral connection
+        assert r is not None, "r needs to be specified!"
+        # edge_index_a2sa = radius(x=pos_s[mask_s_sa, :2], y=pos_s[:, :2], r=r,
+        #                          batch_x=batch_s[mask_s_sa], batch_y=batch_s, max_num_neighbors=max_num_neighbors)
+        edge_index_a2sa = radius(x=pos_s[:, :2], y=pos_s[mask_s_sa, :2], r=r,
+                                 batch_x=batch_s, batch_y=batch_s[mask_s_sa], max_num_neighbors=max_num_neighbors)
+        edge_index_a2sa = edge_index_a2sa[[1, 0]]
+        edge_index_a2sa = edge_index_a2sa[:, ~mask_s_sa[edge_index_a2sa[0]] & mask_s[edge_index_a2sa[0]]]
+
+        # only for seed agent sequence training
+        if seq_mask is not None:
+            edge_mask = seq_mask[edge_index_a2sa[1]]
+            edge_mask = torch.gather(edge_mask, dim=1, index=edge_index_a2sa[0, :, None] % num_agent)[:, 0]
+            edge_index_a2sa = edge_index_a2sa[:, edge_mask]
+
+        if seq_index is None:
+            seq_index = torch.zeros(num_agent, device=pos_a.device).long()
+        if seq_index.dim() == 1:
+            seq_index = seq_index[:, None].repeat(1, num_step)
+        seq_index = seq_index.transpose(0, 1).reshape(-1)
+        assert seq_index.shape[0] == pos_s.shape[0], f"Inconsistent lenght {seq_index.shape[0]} and {pos_s.shape[0]}!"
+
+        # convert to global index
+        all_index = torch.arange(pos_s.shape[0], device=pos_a.device).long()
+        sa_index = all_index[mask_s_sa]
+        edge_index_a2sa[1] = sa_index[edge_index_a2sa[1]]
+
+        # plot edge index TODO: now only support bs=1
+        if os.getenv('PLOT_EDGE_INFERENCE', False) and not is_training:
+            num_agent, num_step, _ = pos_a.shape
+            # plot_interact_edge(edge_index_a2sa, data['scenario_id'], data['batch_size_a'].cpu(), 1, num_step,
+            #                    'interact_a2sa_edge_map', **plot_kwargs)
+            plot_interact_edge(edge_index_a2sa, data['scenario_id'], torch.tensor([num_agent - 1]), 1, num_step,
+                               f"interact_a2sa_edge_map_infer_{plot_kwargs['tag']}", **plot_kwargs)
+
+        rel_pos_a2sa = pos_s[edge_index_a2sa[0]] - pos_s[edge_index_a2sa[1]]
+        rel_head_a2sa = wrap_angle(head_s[edge_index_a2sa[0]] - head_s[edge_index_a2sa[1]])
+
+        r_a2sa = torch.stack(
+            [torch.norm(rel_pos_a2sa[:, :2], p=2, dim=-1),
+            angle_between_2d_vectors(ctr_vector=head_vector_s[edge_index_a2sa[1]], nbr_vector=rel_pos_a2sa[:, :2]),
+            rel_head_a2sa,
+            seq_index[edge_index_a2sa[0]] - seq_index[edge_index_a2sa[1]]], dim=-1)
+        r_a2sa = self.r_a2sa_emb(continuous_inputs=r_a2sa, categorical_embs=None)
+
+        return edge_index_a2sa, r_a2sa
+
+    def _build_map2sa_edge(self, data, pos_a, head_a, head_vector_a, batch_s, batch_pl,
+                           mask_sa, inference_mask=None, r=None, max_num_neighbors=32, mode: Literal['grid', 'heading']='heading'):
+
+        _, num_step, _ = pos_a.shape
+
+        mask_pl2sa = torch.ones_like(mask_sa).bool()
+        if inference_mask is not None:
+            mask_pl2sa = mask_pl2sa & inference_mask
+        mask_pl2sa = mask_pl2sa.transpose(0, 1).reshape(-1)
+        mask_s_sa = mask_sa.transpose(0, 1).reshape(-1)
+
+        pos_s = pos_a.transpose(0, 1).reshape(-1, self.input_dim)
+        head_s = head_a.transpose(0, 1).reshape(-1)
+        head_vector_s = head_vector_a.transpose(0, 1).reshape(-1, 2)
+
+        ori_pos_pl = data['pt_token']['position'][:, :self.input_dim].contiguous()
+        ori_orient_pl = data['pt_token']['orientation'].contiguous()
+        pos_pl = ori_pos_pl.repeat(num_step, 1) # not `repeat_interleave`
+        orient_pl = ori_orient_pl.repeat(num_step)
+
+        # build map2agent directed graph
+        assert r is not None, "r needs to be specified!"
+        # edge_index_pl2sa = radius(x=pos_s[mask_s_sa, :2], y=pos_pl[:, :2], r=r,
+        #                           batch_x=batch_s[mask_s_sa], batch_y=batch_pl, max_num_neighbors=max_num_neighbors)
+        edge_index_pl2sa = radius(x=pos_pl[:, :2], y=pos_s[mask_s_sa, :2], r=r,
+                                  batch_x=batch_pl, batch_y=batch_s[mask_s_sa], max_num_neighbors=max_num_neighbors)
+        edge_index_pl2sa = edge_index_pl2sa[[1, 0]]
+        edge_index_pl2sa = edge_index_pl2sa[:, mask_pl2sa[mask_s_sa][edge_index_pl2sa[1]]]
+
+        # convert to global index
+        all_index = torch.arange(pos_s.shape[0], device=pos_a.device).long()
+        sa_index = all_index[mask_s_sa]
+        edge_index_pl2sa[1] = sa_index[edge_index_pl2sa[1]]
+
+        # plot edge map
+        # if os.getenv('PLOT_EDGE', False):
+        #     plot_map_edge(edge_index_pl2sa, pos_s[:, :2], data, save_path='map2sa_edge_map')
+
+        rel_pos_pl2sa = pos_pl[edge_index_pl2sa[0]] - pos_s[edge_index_pl2sa[1]]
+        rel_orient_pl2sa = wrap_angle(orient_pl[edge_index_pl2sa[0]] - head_s[edge_index_pl2sa[1]])
+
+        r_pl2sa = torch.stack(
+            [torch.norm(rel_pos_pl2sa[:, :2], p=2, dim=-1),
+            angle_between_2d_vectors(ctr_vector=head_vector_s[edge_index_pl2sa[1]], nbr_vector=rel_pos_pl2sa[:, :2]),
+            rel_orient_pl2sa], dim=-1)
+        r_pl2sa = self.r_pt2sa_emb(continuous_inputs=r_pl2sa, categorical_embs=None)
+
+        return edge_index_pl2sa, r_pl2sa
+
+    def _build_sa2sa_edge(self, data, pos_a, head_a, state_a, head_vector_a, batch_s, mask, inference_mask=None, **plot_kwargs):
+
+        num_agent = pos_a.shape[0]
+
+        pos_t = pos_a.transpose(0, 1).reshape(-1, self.input_dim)
+        head_t = head_a.reshape(-1)
+        head_vector_t = head_vector_a.reshape(-1, 2)
+
+        if inference_mask is not None:
+            mask_t = mask.unsqueeze(2) & inference_mask.unsqueeze(1)
+        else:
+            mask_t = mask.unsqueeze(2) & mask.unsqueeze(1)
+
+        edge_index_sa2sa = dense_to_sparse(mask_t)[0]
+        edge_index_sa2sa = edge_index_sa2sa[:, edge_index_sa2sa[1] - edge_index_sa2sa[0] > 0]
+        rel_pos_t = pos_t[edge_index_sa2sa[0]] - pos_t[edge_index_sa2sa[1]]
+        rel_head_t = wrap_angle(head_t[edge_index_sa2sa[0]] - head_t[edge_index_sa2sa[1]])
+
+        r_t = torch.stack(
+            [torch.norm(rel_pos_t[:, :2], p=2, dim=-1),
+             angle_between_2d_vectors(ctr_vector=head_vector_t[edge_index_sa2sa[1]], nbr_vector=rel_pos_t[:, :2]),
+             rel_head_t,
+             edge_index_sa2sa[0] - edge_index_sa2sa[1]], dim=-1)
+        r_sa2sa = self.r_sa2sa_emb(continuous_inputs=r_t, categorical_embs=None)
+
+        return edge_index_sa2sa, r_sa2sa
+
+    def _build_seq(self, device, num_agent, num_step, av_index, sort_indices):
+        """
+        Args:
+            sort_indices (torch.Tensor): shape (num_agent, num_atep)
+        """
+        # sort_indices = sort_indices[:self.num_seed_feature]
+        seq_mask = torch.ones(self.num_seed_feature, num_step, num_agent + self.num_seed_feature, device=device).bool()
+        seq_mask[..., -self.num_seed_feature:] = False
+        for t in range(num_step):
+            for s in range(self.num_seed_feature):
+                seq_mask[s, t, sort_indices[s:, t].flatten().long()] = False
+        if self.seed_attn_to_av:
+            seq_mask[..., av_index] = True
+        seq_mask = seq_mask.transpose(0, 1).reshape(-1, num_agent + self.num_seed_feature)
+
+        seq_index = torch.cat([torch.zeros(num_agent), torch.arange(self.num_seed_feature) + 1]).to(device)
+        seq_index = seq_index[:, None].repeat(1, num_step)
+        for t in range(num_step):
+            for s in range(self.num_seed_feature):
+                seq_index[sort_indices[s : s + 1, t].flatten().long(), t] = s + 1
+        seq_index[av_index] = 0
+
+        return seq_mask, seq_index
+
+    def _build_occ_gt(self, data, seq_mask, pos_rel_index_gt, pos_rel_index_gt_seed, mask_seed,
+                      edge_index=None, mode='edge_index'):
+        """
+        Args:
+            seq_mask (torch.Tensor): shape (num_step * num_seed_feature, num_agent + self.num_seed_feature)
+            pos_rel_index_gt (torch.Tensor): shape (num_agent, num_step)
+            pos_rel_index_gt_seed (torch.Tensor): shape (num_seed, num_step)
+        """
+        num_agent = data['agent']['state_idx'].shape[0] + self.num_seed_feature
+        num_step = data['agent']['state_idx'].shape[1]
+        data['agent']['agent_occ'] = torch.zeros(data.num_graphs * self.num_seed_feature, num_step, self.attr_tokenizer.grid_size,
+                                                            device=data['agent']['state_idx'].device).long()
+        data['agent']['map_occ'] = torch.zeros(data.num_graphs, num_step, self.attr_tokenizer.grid_size,
+                                                            device=data['agent']['state_idx'].device).long()
+
+        if mode == 'edge_index':
+
+            assert edge_index is not None, f"Need edge_index input!"
+            for src_index in torch.unique(edge_index[1]):
+                # decode src
+                src_row = src_index % num_agent - (num_agent - self.num_seed_feature)
+                src_col = src_index // num_agent
+                # decode tgt
+                tgt_indexes = edge_index[0, edge_index[1] == src_index]
+                tgt_rows = tgt_indexes % num_agent
+                tgt_cols = tgt_indexes // num_agent
+                assert tgt_rows.max() < num_agent - self.num_seed_feature, f"Invalid {tgt_rows}"
+                assert torch.unique(tgt_cols).shape[0] == 1 and torch.unique(tgt_cols)[0] == src_col
+                data['agent']['agent_occ'][src_row, src_col, pos_rel_index_gt[tgt_rows, tgt_cols]] = 1
+
+        else:
+
+            seq_mask = seq_mask.reshape(num_step, self.num_seed_feature, -1).transpose(0, 1)[..., :-self.num_seed_feature]
+            for s in range(self.num_seed_feature):
+                for t in range(num_step):
+                    index = pos_rel_index_gt[seq_mask[s, t], t]
+                    data['agent']['agent_occ'][s, t, index[index != -1]] = 1
+                    if t > 0 and s < pos_rel_index_gt_seed.shape[0] and mask_seed[s, t - 1]: # insert agents
+                        data['agent']['agent_occ'][s, t, pos_rel_index_gt_seed[s, t - 1]] = -1
+
+        # TODO: fix batch_size!!!
+        pt_grid_token_idx = data['agent']['pt_grid_token_idx'] # (t, num_pt)
+        for t in range(num_step):
+            data['agent']['map_occ'][:, t, pt_grid_token_idx[t][pt_grid_token_idx[t] != -1]] = 1
+        data['agent']['map_occ'] = data['agent']['map_occ'].repeat_interleave(repeats=self.num_seed_feature, dim=0)
+
+    def forward(self,
+                data: HeteroData,
+                map_enc: Mapping[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
+
+        pos_a = data['agent']['token_pos'].clone() # (a, t, 2)
+        head_a = data['agent']['token_heading'].clone() # (a, t)
+        num_agent, num_step, traj_dim = pos_a.shape # e.g. (50, 18, 2)
+        num_pt = data['pt_token']['position'].shape[0]
+        agent_category = data['agent']['category'].clone() # (a,)
+        agent_shape = data['agent']['shape'][:, self.num_historical_steps - 1].clone() # (a, 3)
+        agent_token_index = data['agent']['token_idx'].clone() # (a, t)
+        agent_state_index = data['agent']['state_idx'].clone()
+        agent_type_index = data['agent']['type'].clone()
+
+        av_index = data['agent']['av_index'].long()
+        ego_pos = pos_a[av_index]
+        ego_head = head_a[av_index]
+
+        _, head_vector_a = self._build_vector_a(pos_a, head_a, agent_state_index)
+
+        agent_grid_token_idx = data['agent']['grid_token_idx']
+        agent_grid_offset_xy = data['agent']['grid_offset_xy']
+        agent_head_token_idx = data['agent']['heading_token_idx']
+        sort_indices = data['agent']['sort_indices']
+        pt_grid_token_idx = data['agent']['pt_grid_token_idx']
+
+        ori_pos_pl = data['pt_token']['position'][:, :self.input_dim].contiguous()
+        ori_orient_pl = data['pt_token']['orientation'].contiguous()
+        pos_pl = ori_pos_pl.repeat(num_step, 1)
+        orient_pl = ori_orient_pl.repeat(num_step)
+
+        # build relative 3d descriptors
+        pos_s = pos_a.transpose(0, 1).reshape(-1, self.input_dim)
+        head_s = head_a.transpose(0, 1).reshape(-1)
+
+        ego_pos_a = ego_pos.repeat_interleave(repeats=data['batch_size_a'], dim=0)
+        ego_head_a = ego_head.repeat_interleave(repeats=data['batch_size_a'], dim=0)
+        ego_pos_s = ego_pos_a.transpose(0, 1).reshape(-1, self.input_dim)
+        ego_head_s = ego_head_a.transpose(0, 1).reshape(-1)
+        rel_pos_a2a = pos_s - ego_pos_s
+        rel_head_a2a = head_s - ego_head_s
+
+        ego_pos_pl = ego_pos.repeat_interleave(repeats=data['batch_size_pl'], dim=0)
+        ego_head_pl = ego_head.repeat_interleave(repeats=data['batch_size_pl'], dim=0)
+        ego_pos_s = ego_pos_pl.transpose(0, 1).reshape(-1, self.input_dim)
+        ego_head_s = ego_head_pl.transpose(0, 1).reshape(-1)
+        rel_pos_pl2a = pos_pl - ego_pos_s
+        rel_head_pl2a = orient_pl - ego_head_s
+
+        # releative encodings
+        ego_head_vector_a = head_vector_a[av_index].repeat_interleave(repeats=data['batch_size_a'], dim=0)
+        ego_head_vector_s = ego_head_vector_a.transpose(0, 1).reshape(-1, 2)
+        r_a2a = torch.stack(
+            [torch.norm(rel_pos_a2a[:, :2], p=2, dim=-1),
+            angle_between_2d_vectors(ctr_vector=ego_head_vector_s, nbr_vector=rel_pos_a2a[:, :2]),
+            rel_head_a2a], dim=-1)
+        r_a2a = self.r_a2a_emb(continuous_inputs=r_a2a, categorical_embs=None)  # [N, hidden_dim]
+
+        ego_head_vector_a = head_vector_a[av_index].repeat_interleave(repeats=data['batch_size_pl'], dim=0)
+        ego_head_vector_s = ego_head_vector_a.transpose(0, 1).reshape(-1, 2)
+        r_pl2a = torch.stack(
+            [torch.norm(rel_pos_pl2a[:, :2], p=2, dim=-1),
+            angle_between_2d_vectors(ctr_vector=ego_head_vector_s, nbr_vector=rel_pos_pl2a[:, :2]),
+            rel_head_pl2a], dim=-1)
+        r_pl2a = self.r_pt2a_emb(continuous_inputs=r_pl2a, categorical_embs=None)  # [M, d]
+
+        r_a2a = r_a2a.reshape(num_step, num_agent, -1).transpose(0, 1)
+        r_pl2a = r_pl2a.reshape(num_step, num_pt, -1).transpose(0, 1)
+        select_agent = torch.randperm(num_agent)[:self.agent_limit]
+        select_pt = torch.randperm(num_pt)[:self.pt_limit]
+        r_a2a = r_a2a[select_agent]
+        r_pl2a = r_pl2a[select_pt]
+
+        # aggregate to global feature
+        r_a2a = r_a2a.mean(dim=0)  # [t, d]
+        r_pl2a = r_pl2a.mean(dim=0)
+
+        # decode grid index of neighbor agents
+        agent_occ = self.grid_agent_occ_head(r_a2a)  # [t, grid_size]
+        pt_occ = self.grid_pt_occ_head(r_pl2a)
+
+        # 1.
+        # agent_occ_gt = torch.zeros_like(agent_occ).long()
+        # pt_occ_gt = torch.zeros_like(pt_occ).long()
+
+        # for t in range(num_step):
+        #     agent_occ_gt[t, agent_grid_token_idx[:, t][agent_grid_token_idx[:, t] != -1]] = 1
+        #     pt_occ_gt[t, pt_grid_token_idx[t][pt_grid_token_idx[t] != -1]] = 1
+
+        # agent_occ_gt[:, self.grid_size // 2] = 0
+        # pt_occ_gt[:, self.grid_size // 2] = 0
+
+        # agent_occ_eval_mask = torch.ones_like(agent_occ_gt)
+        # agent_occ_eval_mask[0] = 0
+        # agent_occ_eval_mask[:, self.grid_size // 2] = 0
+        # pt_occ_eval_mask = torch.ones_like(pt_occ_gt)
+        # pt_occ_eval_mask[0] = 0
+        # pt_occ_eval_mask[:, self.grid_size // 2] = 0
+
+        # 2.
+        # agent_occ_gt = agent_grid_token_idx.transpose(0, 1).reshape(-1)
+        # pt_occ_gt = pt_grid_token_idx.reshape(-1)
+
+        # agent_occ_eval_mask = torch.zeros_like(agent_occ_gt)
+        # agent_occ_eval_mask[torch.randperm(agent_occ_gt.shape[0])[:(num_step * 10)]] = 1
+        # agent_occ_eval_mask[agent_occ_gt == -1] = 0
+
+        # pt_occ_eval_mask = torch.zeros_like(pt_occ_gt)
+        # pt_occ_eval_mask[torch.randperm(pt_occ_gt.shape[0])[:(num_step * 300)]] = 1
+        # pt_occ_eval_mask[pt_occ_gt == -1] = 0
+
+        # 3.
+        agent_occ = agent_occ.reshape(num_step, self.agent_limit, -1)
+        pt_occ = pt_occ.reshape(num_step, self.pt_limit, -1)
+        agent_occ_gt = agent_grid_token_idx[select_agent].transpose(0, 1)
+        pt_occ_gt = pt_grid_token_idx[:, select_pt]
+        agent_occ_eval_mask = agent_occ_gt != -1
+        pt_occ_eval_mask = pt_occ_gt != -1
+
+        agent_occ = agent_occ[:, :agent_occ_gt.shape[1]]
+        pt_occ = pt_occ[:, :pt_occ_gt.shape[1]]
+
+        return {'occ_decoder': True,
+                'num_step': num_step,
+                'num_agent': self.agent_limit, # num_agent
+                'num_pt': self.pt_limit, # num_pt
+                'agent_occ': agent_occ,
+                'agent_occ_gt': agent_occ_gt,
+                'agent_occ_eval_mask': agent_occ_eval_mask.bool(),
+                'pt_occ': pt_occ,
+                'pt_occ_gt': pt_occ_gt,
+                'pt_occ_eval_mask': pt_occ_eval_mask.bool(),
+                }
+
+    def inference(self, *args, **kwargs):
+        return self(*args, **kwargs)
+

seq_10_150_3_3_encode_occ_separate_offsets_bs8_128_no_seqindex_long_lastvalid_ablation_grid/backups/dev/modules/smart_decoder.py

@@ -0,0 +1,137 @@
+from typing import Dict, Optional
+import torch
+import torch.nn as nn
+from torch_geometric.data import HeteroData
+from dev.modules.attr_tokenizer import Attr_Tokenizer
+from dev.modules.agent_decoder import SMARTAgentDecoder
+from dev.modules.occ_decoder import SMARTOccDecoder
+from dev.modules.map_decoder import SMARTMapDecoder
+
+
+DECODER = {'agent_decoder': SMARTAgentDecoder,
+           'occ_decoder': SMARTOccDecoder}
+
+
+class SMARTDecoder(nn.Module):
+
+    def __init__(self,
+                 decoder_type: str,
+                 dataset: str,
+                 input_dim: int,
+                 hidden_dim: int,
+                 num_historical_steps: int,
+                 pl2pl_radius: float,
+                 time_span: Optional[int],
+                 pl2a_radius: float,
+                 pl2seed_radius: float,
+                 a2a_radius: float,
+                 a2sa_radius: float,
+                 pl2sa_radius: float,
+                 num_freq_bands: int,
+                 num_map_layers: int,
+                 num_agent_layers: int,
+                 num_heads: int,
+                 head_dim: int,
+                 dropout: float,
+                 map_token: Dict,
+                 token_size=512,
+                 attr_tokenizer: Attr_Tokenizer=None,
+                 predict_motion: bool=False,
+                 predict_state: bool=False,
+                 predict_map: bool=False,
+                 predict_occ: bool=False,
+                 use_grid_token: bool=False,
+                 state_token: Dict[str, int]=None,
+                 seed_size: int=5,
+                 buffer_size: int=32,
+                 num_recurrent_steps_val: int=-1,
+                 loss_weight: dict=None,
+                 logger=None) -> None:
+
+        super(SMARTDecoder, self).__init__()
+
+        self.map_encoder = SMARTMapDecoder(
+            dataset=dataset,
+            input_dim=input_dim,
+            hidden_dim=hidden_dim,
+            num_historical_steps=num_historical_steps,
+            pl2pl_radius=pl2pl_radius,
+            num_freq_bands=num_freq_bands,
+            num_layers=num_map_layers,
+            num_heads=num_heads,
+            head_dim=head_dim,
+            dropout=dropout,
+            map_token=map_token,
+        )
+
+        assert decoder_type in list(DECODER.keys()), f"Unsupport decoder type: {decoder_type}"
+        self.agent_encoder = DECODER[decoder_type](
+            dataset=dataset,
+            input_dim=input_dim,
+            hidden_dim=hidden_dim,
+            num_historical_steps=num_historical_steps,
+            time_span=time_span,
+            pl2a_radius=pl2a_radius,
+            pl2seed_radius=pl2seed_radius,
+            a2a_radius=a2a_radius,
+            a2sa_radius=a2sa_radius,
+            pl2sa_radius=pl2sa_radius,
+            num_freq_bands=num_freq_bands,
+            num_layers=num_agent_layers,
+            num_heads=num_heads,
+            head_dim=head_dim,
+            dropout=dropout,
+            token_size=token_size,
+            attr_tokenizer=attr_tokenizer,
+            predict_motion=predict_motion,
+            predict_state=predict_state,
+            predict_map=predict_map,
+            predict_occ=predict_occ,
+            state_token=state_token,
+            use_grid_token=use_grid_token,
+            seed_size=seed_size,
+            buffer_size=buffer_size,
+            num_recurrent_steps_val=num_recurrent_steps_val,
+            loss_weight=loss_weight,
+            logger=logger,
+        )
+        self.map_enc = None
+        self.predict_motion = predict_motion
+        self.predict_state = predict_state
+        self.predict_map = predict_map
+        self.predict_occ = predict_occ
+        self.data_keys = ["agent_valid_mask", "category", "valid_mask", "av_index", "scenario_id", "shape"]
+
+    def get_agent_inputs(self, data: HeteroData) -> Dict[str, torch.Tensor]:
+        return self.agent_encoder.get_inputs(data)
+
+    def forward(self, data: HeteroData) -> Dict[str, torch.Tensor]:
+        map_enc = self.map_encoder(data)
+
+        agent_enc = {}
+        if self.predict_motion or self.predict_state or self.predict_occ:
+            agent_enc = self.agent_encoder(data, map_enc)
+
+        return {**map_enc, **agent_enc, **{k: data[k] for k in self.data_keys}}
+
+    def inference(self, data: HeteroData) -> Dict[str, torch.Tensor]:
+        map_enc = self.map_encoder(data)
+
+        agent_enc = {}
+        if self.predict_motion or self.predict_state or self.predict_occ:
+            agent_enc = self.agent_encoder.inference(data, map_enc)
+
+        return {**map_enc, **agent_enc, **{k: data[k] for k in self.data_keys}}
+
+    def inference_no_map(self, data: HeteroData, map_enc) -> Dict[str, torch.Tensor]:
+        agent_enc = self.agent_encoder.inference(data, map_enc)
+        return {**map_enc, **agent_enc}
+
+    def insert_agent(self, data: HeteroData) -> Dict[str, torch.Tensor]:
+        map_enc = self.map_encoder(data)
+        agent_enc = self.agent_encoder.insert(data, map_enc)
+        return {**map_enc, **agent_enc, **{k: data[k] for k in self.data_keys}}
+
+    def predict_nearest_pos(self, data: HeteroData, rank) -> Dict[str, torch.Tensor]:
+        map_enc = self.map_encoder(data)
+        self.agent_encoder.predict_nearest_pos(data, map_enc, rank)

seq_10_150_3_3_encode_occ_separate_offsets_bs8_128_no_seqindex_long_lastvalid_ablation_grid/backups/dev/utils/cluster_reader.py

@@ -0,0 +1,45 @@
+import io
+import pickle
+import pandas as pd
+import json
+
+
+class LoadScenarioFromCeph:
+    def __init__(self):
+        from petrel_client.client import Client
+        self.file_client = Client('~/petreloss.conf')
+
+    def list(self, dir_path):
+        return list(self.file_client.list(dir_path))
+
+    def save(self, data, url):
+        self.file_client.put(url, pickle.dumps(data))
+
+    def read_correct_csv(self, scenario_path):
+        output = pd.read_csv(io.StringIO(self.file_client.get(scenario_path).decode('utf-8')), engine="python")
+        return output
+
+    def contains(self, url):
+        return self.file_client.contains(url)
+
+    def read_string(self, csv_url):
+        from io import StringIO
+        df = pd.read_csv(StringIO(str(self.file_client.get(csv_url), 'utf-8')), sep='\s+', low_memory=False)
+        return df
+
+    def read(self, scenario_path):
+        with io.BytesIO(self.file_client.get(scenario_path)) as f:
+            datas = pickle.load(f)
+            return datas
+
+    def read_json(self, path):
+        with io.BytesIO(self.file_client.get(path)) as f:
+            data = json.load(f)
+            return data
+
+    def read_csv(self, scenario_path):
+        return pickle.loads(self.file_client.get(scenario_path))
+
+    def read_model(self, model_path):
+        with io.BytesIO(self.file_client.get(model_path)) as f:
+            pass

seq_10_150_3_3_encode_occ_separate_offsets_bs8_128_no_seqindex_long_lastvalid_ablation_grid/backups/dev/utils/func.py

@@ -0,0 +1,260 @@
+import logging
+import time
+import os
+import yaml
+import easydict
+import math
+import torch
+import torch.nn as nn
+from rich.console import Console
+from typing import Any, List, Optional, Mapping
+
+from lightning_utilities.core.rank_zero import rank_prefixed_message, rank_zero_only
+
+
+CONSOLE = Console(width=128)
+
+
+def check_nan_inf(t, s):
+    assert not torch.isinf(t).any(), f"{s} is inf, {t}"
+    assert not torch.isnan(t).any(), f"{s} is nan, {t}"
+
+
+def safe_list_index(ls: List[Any], elem: Any) -> Optional[int]:
+    try:
+        return ls.index(elem)
+    except ValueError:
+        return None
+
+
+def angle_between_2d_vectors(
+        ctr_vector: torch.Tensor,
+        nbr_vector: torch.Tensor) -> torch.Tensor:
+    return torch.atan2(ctr_vector[..., 0] * nbr_vector[..., 1] - ctr_vector[..., 1] * nbr_vector[..., 0],
+                       (ctr_vector[..., :2] * nbr_vector[..., :2]).sum(dim=-1))
+
+
+def angle_between_3d_vectors(
+        ctr_vector: torch.Tensor,
+        nbr_vector: torch.Tensor) -> torch.Tensor:
+    return torch.atan2(torch.cross(ctr_vector, nbr_vector, dim=-1).norm(p=2, dim=-1),
+                       (ctr_vector * nbr_vector).sum(dim=-1))
+
+
+def side_to_directed_lineseg(
+        query_point: torch.Tensor,
+        start_point: torch.Tensor,
+        end_point: torch.Tensor) -> str:
+    cond = ((end_point[0] - start_point[0]) * (query_point[1] - start_point[1]) -
+            (end_point[1] - start_point[1]) * (query_point[0] - start_point[0]))
+    if cond > 0:
+        return 'LEFT'
+    elif cond < 0:
+        return 'RIGHT'
+    else:
+        return 'CENTER'
+
+
+def wrap_angle(
+        angle: torch.Tensor,
+        min_val: float = -math.pi,
+        max_val: float = math.pi) -> torch.Tensor:
+    return min_val + (angle + max_val) % (max_val - min_val)
+
+
+def load_config_act(path):
+    """ load config file"""
+    with open(path, 'r') as f:
+        cfg = yaml.load(f, Loader=yaml.FullLoader)
+    return easydict.EasyDict(cfg)
+
+
+def load_config_init(path):
+    """ load config file"""
+    path = os.path.join('init/configs', f'{path}.yaml')
+    with open(path, 'r') as f:
+        cfg = yaml.load(f, Loader=yaml.FullLoader)
+    return cfg
+
+
+class Logging:
+
+    def make_log_dir(self, dirname='logs'):
+        now_dir = os.path.dirname(__file__)
+        path = os.path.join(now_dir, dirname)
+        path = os.path.normpath(path)
+        if not os.path.exists(path):
+            os.mkdir(path)
+        return path
+
+    def get_log_filename(self):
+        filename = "{}.log".format(time.strftime("%Y-%m-%d-%H%M%S", time.localtime()))
+        filename = os.path.join(self.make_log_dir(), filename)
+        filename = os.path.normpath(filename)
+        return filename
+
+    def log(self, level='DEBUG', name="simagent"):
+        logger = logging.getLogger(name)
+        level = getattr(logging, level)
+        logger.setLevel(level)
+        if not logger.handlers:
+            sh = logging.StreamHandler()
+            fh = logging.FileHandler(filename=self.get_log_filename(), mode='a',encoding="utf-8")
+            fmt = logging.Formatter("%(asctime)s-%(levelname)s-%(filename)s-Line:%(lineno)d-Message:%(message)s")
+            sh.setFormatter(fmt=fmt)
+            fh.setFormatter(fmt=fmt)
+            logger.addHandler(sh)
+            logger.addHandler(fh)
+        return logger
+
+    def add_log(self, logger, level='DEBUG'):
+        level = getattr(logging, level)
+        logger.setLevel(level)
+        if not logger.handlers:
+            sh = logging.StreamHandler()
+            fh = logging.FileHandler(filename=self.get_log_filename(), mode='a',encoding="utf-8")
+            fmt = logging.Formatter("%(asctime)s-%(levelname)s-%(filename)s-Line:%(lineno)d-Message:%(message)s")
+            sh.setFormatter(fmt=fmt)
+            fh.setFormatter(fmt=fmt)
+            logger.addHandler(sh)
+            logger.addHandler(fh)
+        return logger
+
+
+# Adapted from 'CatK'
+class RankedLogger(logging.LoggerAdapter):
+    """A multi-GPU-friendly python command line logger."""
+
+    def __init__(
+        self,
+        name: str = __name__,
+        rank_zero_only: bool = False,
+        extra: Optional[Mapping[str, object]] = None,
+    ) -> None:
+        """Initializes a multi-GPU-friendly python command line logger that logs on all processes
+        with their rank prefixed in the log message.
+
+        :param name: The name of the logger. Default is ``__name__``.
+        :param rank_zero_only: Whether to force all logs to only occur on the rank zero process. Default is `False`.
+        :param extra: (Optional) A dict-like object which provides contextual information. See `logging.LoggerAdapter`.
+        """
+        logger = logging.getLogger(name)
+        super().__init__(logger=logger, extra=extra)
+        self.rank_zero_only = rank_zero_only
+
+    def log(
+        self, level: int, msg: str, rank: Optional[int] = None, *args, **kwargs
+    ) -> None:
+        """Delegate a log call to the underlying logger, after prefixing its message with the rank
+        of the process it's being logged from. If `'rank'` is provided, then the log will only
+        occur on that rank/process.
+
+        :param level: The level to log at. Look at `logging.__init__.py` for more information.
+        :param msg: The message to log.
+        :param rank: The rank to log at.
+        :param args: Additional args to pass to the underlying logging function.
+        :param kwargs: Any additional keyword args to pass to the underlying logging function.
+        """
+        if self.isEnabledFor(level):
+            msg, kwargs = self.process(msg, kwargs)
+            current_rank = getattr(rank_zero_only, "rank", None)
+            if current_rank is None:
+                raise RuntimeError(
+                    "The `rank_zero_only.rank` needs to be set before use"
+                )
+            msg = rank_prefixed_message(msg, current_rank)
+            if self.rank_zero_only:
+                if current_rank == 0:
+                    self.logger.log(level, msg, *args, **kwargs)
+            else:
+                if rank is None:
+                    self.logger.log(level, msg, *args, **kwargs)
+                elif current_rank == rank:
+                    self.logger.log(level, msg, *args, **kwargs)
+
+
+
+def weight_init(m: nn.Module) -> None:
+    if isinstance(m, nn.Linear):
+        nn.init.xavier_uniform_(m.weight)
+        if m.bias is not None:
+            nn.init.zeros_(m.bias)
+    elif isinstance(m, (nn.Conv1d, nn.Conv2d, nn.Conv3d)):
+        fan_in = m.in_channels / m.groups
+        fan_out = m.out_channels / m.groups
+        bound = (6.0 / (fan_in + fan_out)) ** 0.5
+        nn.init.uniform_(m.weight, -bound, bound)
+        if m.bias is not None:
+            nn.init.zeros_(m.bias)
+    elif isinstance(m, nn.Embedding):
+        nn.init.normal_(m.weight, mean=0.0, std=0.02)
+    elif isinstance(m, (nn.BatchNorm1d, nn.BatchNorm2d, nn.BatchNorm3d)):
+        nn.init.ones_(m.weight)
+        nn.init.zeros_(m.bias)
+    elif isinstance(m, nn.LayerNorm):
+        nn.init.ones_(m.weight)
+        nn.init.zeros_(m.bias)
+    elif isinstance(m, nn.MultiheadAttention):
+        if m.in_proj_weight is not None:
+            fan_in = m.embed_dim
+            fan_out = m.embed_dim
+            bound = (6.0 / (fan_in + fan_out)) ** 0.5
+            nn.init.uniform_(m.in_proj_weight, -bound, bound)
+        else:
+            nn.init.xavier_uniform_(m.q_proj_weight)
+            nn.init.xavier_uniform_(m.k_proj_weight)
+            nn.init.xavier_uniform_(m.v_proj_weight)
+        if m.in_proj_bias is not None:
+            nn.init.zeros_(m.in_proj_bias)
+        nn.init.xavier_uniform_(m.out_proj.weight)
+        if m.out_proj.bias is not None:
+            nn.init.zeros_(m.out_proj.bias)
+        if m.bias_k is not None:
+            nn.init.normal_(m.bias_k, mean=0.0, std=0.02)
+        if m.bias_v is not None:
+            nn.init.normal_(m.bias_v, mean=0.0, std=0.02)
+    elif isinstance(m, (nn.LSTM, nn.LSTMCell)):
+        for name, param in m.named_parameters():
+            if 'weight_ih' in name:
+                for ih in param.chunk(4, 0):
+                    nn.init.xavier_uniform_(ih)
+            elif 'weight_hh' in name:
+                for hh in param.chunk(4, 0):
+                    nn.init.orthogonal_(hh)
+            elif 'weight_hr' in name:
+                nn.init.xavier_uniform_(param)
+            elif 'bias_ih' in name:
+                nn.init.zeros_(param)
+            elif 'bias_hh' in name:
+                nn.init.zeros_(param)
+                nn.init.ones_(param.chunk(4, 0)[1])
+    elif isinstance(m, (nn.GRU, nn.GRUCell)):
+        for name, param in m.named_parameters():
+            if 'weight_ih' in name:
+                for ih in param.chunk(3, 0):
+                    nn.init.xavier_uniform_(ih)
+            elif 'weight_hh' in name:
+                for hh in param.chunk(3, 0):
+                    nn.init.orthogonal_(hh)
+            elif 'bias_ih' in name:
+                nn.init.zeros_(param)
+            elif 'bias_hh' in name:
+                nn.init.zeros_(param)
+
+
+def pos2posemb(pos, num_pos_feats=128, temperature=10000):
+
+    scale = 2 * math.pi
+    pos = pos * scale
+    dim_t = torch.arange(num_pos_feats, dtype=torch.float32, device=pos.device)
+    dim_t = temperature ** (2 * (dim_t // 2) / num_pos_feats)
+
+    D = pos.shape[-1]
+    pos_dims = []
+    for i in range(D):
+        pos_dim_i = pos[..., i, None] / dim_t
+        pos_dim_i = torch.stack((pos_dim_i[..., 0::2].sin(), pos_dim_i[..., 1::2].cos()), dim=-1).flatten(-2)
+        pos_dims.append(pos_dim_i)
+        posemb = torch.cat(pos_dims, dim=-1)
+
+    return posemb

seq_10_150_3_3_encode_occ_separate_offsets_bs8_128_no_seqindex_long_lastvalid_ablation_grid/backups/dev/utils/graph.py

@@ -0,0 +1,89 @@
+
+import torch
+from typing import List, Optional, Tuple, Union
+from torch_geometric.utils import coalesce
+from torch_geometric.utils import degree
+
+
+def add_edges(
+        from_edge_index: torch.Tensor,
+        to_edge_index: torch.Tensor,
+        from_edge_attr: Optional[torch.Tensor] = None,
+        to_edge_attr: Optional[torch.Tensor] = None,
+        replace: bool = True) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
+    from_edge_index = from_edge_index.to(device=to_edge_index.device, dtype=to_edge_index.dtype)
+    mask = ((to_edge_index[0].unsqueeze(-1) == from_edge_index[0].unsqueeze(0)) &
+            (to_edge_index[1].unsqueeze(-1) == from_edge_index[1].unsqueeze(0)))
+    if replace:
+        to_mask = mask.any(dim=1)
+        if from_edge_attr is not None and to_edge_attr is not None:
+            from_edge_attr = from_edge_attr.to(device=to_edge_attr.device, dtype=to_edge_attr.dtype)
+            to_edge_attr = torch.cat([to_edge_attr[~to_mask], from_edge_attr], dim=0)
+        to_edge_index = torch.cat([to_edge_index[:, ~to_mask], from_edge_index], dim=1)
+    else:
+        from_mask = mask.any(dim=0)
+        if from_edge_attr is not None and to_edge_attr is not None:
+            from_edge_attr = from_edge_attr.to(device=to_edge_attr.device, dtype=to_edge_attr.dtype)
+            to_edge_attr = torch.cat([to_edge_attr, from_edge_attr[~from_mask]], dim=0)
+        to_edge_index = torch.cat([to_edge_index, from_edge_index[:, ~from_mask]], dim=1)
+    return to_edge_index, to_edge_attr
+
+
+def merge_edges(
+        edge_indices: List[torch.Tensor],
+        edge_attrs: Optional[List[torch.Tensor]] = None,
+        reduce: str = 'add') -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
+    edge_index = torch.cat(edge_indices, dim=1)
+    if edge_attrs is not None:
+        edge_attr = torch.cat(edge_attrs, dim=0)
+    else:
+        edge_attr = None
+    return coalesce(edge_index=edge_index, edge_attr=edge_attr, reduce=reduce)
+
+
+def complete_graph(
+        num_nodes: Union[int, Tuple[int, int]],
+        ptr: Optional[Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]] = None,
+        loop: bool = False,
+        device: Optional[Union[torch.device, str]] = None) -> torch.Tensor:
+    if ptr is None:
+        if isinstance(num_nodes, int):
+            num_src, num_dst = num_nodes, num_nodes
+        else:
+            num_src, num_dst = num_nodes
+        edge_index = torch.cartesian_prod(torch.arange(num_src, dtype=torch.long, device=device),
+                                          torch.arange(num_dst, dtype=torch.long, device=device)).t()
+    else:
+        if isinstance(ptr, torch.Tensor):
+            ptr_src, ptr_dst = ptr, ptr
+            num_src_batch = num_dst_batch = ptr[1:] - ptr[:-1]
+        else:
+            ptr_src, ptr_dst = ptr
+            num_src_batch = ptr_src[1:] - ptr_src[:-1]
+            num_dst_batch = ptr_dst[1:] - ptr_dst[:-1]
+        edge_index = torch.cat(
+            [torch.cartesian_prod(torch.arange(num_src, dtype=torch.long, device=device),
+                                  torch.arange(num_dst, dtype=torch.long, device=device)) + p
+             for num_src, num_dst, p in zip(num_src_batch, num_dst_batch, torch.stack([ptr_src, ptr_dst], dim=1))],
+            dim=0)
+        edge_index = edge_index.t()
+    if isinstance(num_nodes, int) and not loop:
+        edge_index = edge_index[:, edge_index[0] != edge_index[1]]
+    return edge_index.contiguous()
+
+
+def bipartite_dense_to_sparse(adj: torch.Tensor) -> torch.Tensor:
+    index = adj.nonzero(as_tuple=True)
+    if len(index) == 3:
+        batch_src = index[0] * adj.size(1)
+        batch_dst = index[0] * adj.size(2)
+        index = (batch_src + index[1], batch_dst + index[2])
+    return torch.stack(index, dim=0)
+
+
+def unbatch(
+        src: torch.Tensor,
+        batch: torch.Tensor,
+        dim: int = 0) -> List[torch.Tensor]:
+    sizes = degree(batch, dtype=torch.long).tolist()
+    return src.split(sizes, dim)

seq_10_150_3_3_encode_occ_separate_offsets_bs8_128_no_seqindex_long_lastvalid_ablation_grid/backups/dev/utils/metrics.py

@@ -0,0 +1,692 @@
+import torch
+import os
+import itertools
+import multiprocessing as mp
+import torch.nn.functional as F
+from pathlib import Path
+from torch.nn import CrossEntropyLoss
+from torch_scatter import gather_csr
+from torch_scatter import segment_csr
+from torchmetrics import Metric
+from typing import Optional, Tuple, Dict, List
+
+
+__all__ = ['minADE', 'minFDE', 'TokenCls', 'StateAccuracy', 'GridOverlapRate']
+
+
+class CustomCrossEntropyLoss(CrossEntropyLoss):
+
+    def __init__(self, label_smoothing=0.0, reduction='mean'):
+        super(CustomCrossEntropyLoss, self).__init__()
+        self.label_smoothing = label_smoothing
+        self.reduction = reduction
+
+    def forward(self, input, target):
+        num_classes = input.size(1)
+
+        log_probs = F.log_softmax(input, dim=1)
+        
+        with torch.no_grad():
+            smooth_target = torch.zeros_like(log_probs).scatter_(1, target.unsqueeze(1), 1)
+            smooth_target = smooth_target * (1 - self.label_smoothing) + self.label_smoothing / num_classes
+        
+        loss = -torch.sum(log_probs * smooth_target, dim=1)
+        
+        if self.reduction == 'mean':
+            return loss.mean()
+        elif self.reduction == 'sum':
+            return loss.sum()
+        else:
+            return loss
+
+
+def topk(
+        max_guesses: int,
+        pred: torch.Tensor,
+        prob: Optional[torch.Tensor] = None,
+        ptr: Optional[torch.Tensor] = None,
+        joint: bool = False) -> Tuple[torch.Tensor, torch.Tensor]:
+    max_guesses = min(max_guesses, pred.size(1))
+    if max_guesses == pred.size(1):
+        if prob is not None:
+            prob = prob / prob.sum(dim=-1, keepdim=True)
+        else:
+            prob = pred.new_ones((pred.size(0), max_guesses)) / max_guesses
+        return pred, prob
+    else:
+        if prob is not None:
+            if joint:
+                if ptr is None:
+                    inds_topk = torch.topk((prob / prob.sum(dim=-1, keepdim=True)).mean(dim=0, keepdim=True),
+                                           k=max_guesses, dim=-1, largest=True, sorted=True)[1]
+                    inds_topk = inds_topk.repeat(pred.size(0), 1)
+                else:
+                    inds_topk = torch.topk(segment_csr(src=prob / prob.sum(dim=-1, keepdim=True), indptr=ptr,
+                                                       reduce='mean'),
+                                           k=max_guesses, dim=-1, largest=True, sorted=True)[1]
+                    inds_topk = gather_csr(src=inds_topk, indptr=ptr)
+            else:
+                inds_topk = torch.topk(prob, k=max_guesses, dim=-1, largest=True, sorted=True)[1]
+            pred_topk = pred[torch.arange(pred.size(0)).unsqueeze(-1).expand(-1, max_guesses), inds_topk]
+            prob_topk = prob[torch.arange(pred.size(0)).unsqueeze(-1).expand(-1, max_guesses), inds_topk]
+            prob_topk = prob_topk / prob_topk.sum(dim=-1, keepdim=True)
+        else:
+            pred_topk = pred[:, :max_guesses]
+            prob_topk = pred.new_ones((pred.size(0), max_guesses)) / max_guesses
+        return pred_topk, prob_topk
+
+
+def topkind(
+        max_guesses: int,
+        pred: torch.Tensor,
+        prob: Optional[torch.Tensor] = None,
+        ptr: Optional[torch.Tensor] = None,
+        joint: bool = False) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    max_guesses = min(max_guesses, pred.size(1))
+    if max_guesses == pred.size(1):
+        if prob is not None:
+            prob = prob / prob.sum(dim=-1, keepdim=True)
+        else:
+            prob = pred.new_ones((pred.size(0), max_guesses)) / max_guesses
+        return pred, prob, None
+    else:
+        if prob is not None:
+            if joint:
+                if ptr is None:
+                    inds_topk = torch.topk((prob / prob.sum(dim=-1, keepdim=True)).mean(dim=0, keepdim=True),
+                                           k=max_guesses, dim=-1, largest=True, sorted=True)[1]
+                    inds_topk = inds_topk.repeat(pred.size(0), 1)
+                else:
+                    inds_topk = torch.topk(segment_csr(src=prob / prob.sum(dim=-1, keepdim=True), indptr=ptr,
+                                                       reduce='mean'),
+                                           k=max_guesses, dim=-1, largest=True, sorted=True)[1]
+                    inds_topk = gather_csr(src=inds_topk, indptr=ptr)
+            else:
+                inds_topk = torch.topk(prob, k=max_guesses, dim=-1, largest=True, sorted=True)[1]
+            pred_topk = pred[torch.arange(pred.size(0)).unsqueeze(-1).expand(-1, max_guesses), inds_topk]
+            prob_topk = prob[torch.arange(pred.size(0)).unsqueeze(-1).expand(-1, max_guesses), inds_topk]
+            prob_topk = prob_topk / prob_topk.sum(dim=-1, keepdim=True)
+        else:
+            pred_topk = pred[:, :max_guesses]
+            prob_topk = pred.new_ones((pred.size(0), max_guesses)) / max_guesses
+        return pred_topk, prob_topk, inds_topk
+
+
+def valid_filter(
+        pred: torch.Tensor,
+        target: torch.Tensor,
+        prob: Optional[torch.Tensor] = None,
+        valid_mask: Optional[torch.Tensor] = None,
+        ptr: Optional[torch.Tensor] = None,
+        keep_invalid_final_step: bool = True) -> Tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor],
+                                                       torch.Tensor, torch.Tensor]:
+    if valid_mask is None:
+        valid_mask = target.new_ones(target.size()[:-1], dtype=torch.bool)
+    if keep_invalid_final_step:
+        filter_mask = valid_mask.any(dim=-1)
+    else:
+        filter_mask = valid_mask[:, -1]
+    pred = pred[filter_mask]
+    target = target[filter_mask]
+    if prob is not None:
+        prob = prob[filter_mask]
+    valid_mask = valid_mask[filter_mask]
+    if ptr is not None:
+        num_nodes_batch = segment_csr(src=filter_mask.long(), indptr=ptr, reduce='sum')
+        ptr = num_nodes_batch.new_zeros((num_nodes_batch.size(0) + 1,))
+        torch.cumsum(num_nodes_batch, dim=0, out=ptr[1:])
+    else:
+        ptr = target.new_tensor([0, target.size(0)])
+    return pred, target, prob, valid_mask, ptr
+
+
+def new_batch_nms(pred_trajs, dist_thresh, num_ret_modes=6):
+    """
+
+    Args:
+        pred_trajs (batch_size, num_modes, num_timestamps, 7)
+        pred_scores (batch_size, num_modes):
+        dist_thresh (float):
+        num_ret_modes (int, optional): Defaults to 6.
+
+    Returns:
+        ret_trajs (batch_size, num_ret_modes, num_timestamps, 5)
+        ret_scores (batch_size, num_ret_modes)
+        ret_idxs (batch_size, num_ret_modes)
+    """
+    batch_size, num_modes, num_timestamps, num_feat_dim = pred_trajs.shape
+    pred_goals = pred_trajs[:, :, -1, :]
+    dist = (pred_goals[:, :, None, 0:2] - pred_goals[:, None, :, 0:2]).norm(dim=-1)
+    nearby_neighbor = dist < dist_thresh
+    pred_scores = nearby_neighbor.sum(dim=-1) / num_modes
+
+    sorted_idxs = pred_scores.argsort(dim=-1, descending=True)
+    bs_idxs_full = torch.arange(batch_size).type_as(sorted_idxs)[:, None].repeat(1, num_modes)
+    sorted_pred_scores = pred_scores[bs_idxs_full, sorted_idxs]
+    sorted_pred_trajs = pred_trajs[bs_idxs_full, sorted_idxs]  # (batch_size, num_modes, num_timestamps, 7)
+    sorted_pred_goals = sorted_pred_trajs[:, :, -1, :]  # (batch_size, num_modes, 7)
+
+    dist = (sorted_pred_goals[:, :, None, 0:2] - sorted_pred_goals[:, None, :, 0:2]).norm(dim=-1)
+    point_cover_mask = (dist < dist_thresh)
+
+    point_val = sorted_pred_scores.clone()  # (batch_size, N)
+    point_val_selected = torch.zeros_like(point_val)  # (batch_size, N)
+
+    ret_idxs = sorted_idxs.new_zeros(batch_size, num_ret_modes).long()
+    ret_trajs = sorted_pred_trajs.new_zeros(batch_size, num_ret_modes, num_timestamps, num_feat_dim)
+    ret_scores = sorted_pred_trajs.new_zeros(batch_size, num_ret_modes)
+    bs_idxs = torch.arange(batch_size).type_as(ret_idxs)
+
+    for k in range(num_ret_modes):
+        cur_idx = point_val.argmax(dim=-1)  # (batch_size)
+        ret_idxs[:, k] = cur_idx
+
+        new_cover_mask = point_cover_mask[bs_idxs, cur_idx]  # (batch_size, N)
+        point_val = point_val * (~new_cover_mask).float()  # (batch_size, N)
+        point_val_selected[bs_idxs, cur_idx] = -1
+        point_val += point_val_selected
+
+        ret_trajs[:, k] = sorted_pred_trajs[bs_idxs, cur_idx]
+        ret_scores[:, k] = sorted_pred_scores[bs_idxs, cur_idx]
+
+    bs_idxs = torch.arange(batch_size).type_as(sorted_idxs)[:, None].repeat(1, num_ret_modes)
+
+    ret_idxs = sorted_idxs[bs_idxs, ret_idxs]
+    return ret_trajs, ret_scores, ret_idxs
+
+
+def batch_nms(pred_trajs, pred_scores,
+              dist_thresh, num_ret_modes=6,
+              mode='static', speed=None):
+    """
+
+    Args:
+        pred_trajs (batch_size, num_modes, num_timestamps, 7)
+        pred_scores (batch_size, num_modes):
+        dist_thresh (float):
+        num_ret_modes (int, optional): Defaults to 6.
+
+    Returns:
+        ret_trajs (batch_size, num_ret_modes, num_timestamps, 5)
+        ret_scores (batch_size, num_ret_modes)
+        ret_idxs (batch_size, num_ret_modes)
+    """
+    batch_size, num_modes, num_timestamps, num_feat_dim = pred_trajs.shape
+
+    sorted_idxs = pred_scores.argsort(dim=-1, descending=True)
+    bs_idxs_full = torch.arange(batch_size).type_as(sorted_idxs)[:, None].repeat(1, num_modes)
+    sorted_pred_scores = pred_scores[bs_idxs_full, sorted_idxs]
+    sorted_pred_trajs = pred_trajs[bs_idxs_full, sorted_idxs]  # (batch_size, num_modes, num_timestamps, 7)
+    sorted_pred_goals = sorted_pred_trajs[:, :, -1, :]  # (batch_size, num_modes, 7)
+
+    if mode == "speed":
+        scale = torch.ones(batch_size).to(sorted_pred_goals.device)
+        lon_dist_thresh = 4 * scale
+        lat_dist_thresh = 0.5 * scale
+        lon_dist = (sorted_pred_goals[:, :, None, [0]] - sorted_pred_goals[:, None, :, [0]]).norm(dim=-1)
+        lat_dist = (sorted_pred_goals[:, :, None, [1]] - sorted_pred_goals[:, None, :, [1]]).norm(dim=-1)
+        point_cover_mask = (lon_dist < lon_dist_thresh[:, None, None]) & (lat_dist < lat_dist_thresh[:, None, None])
+    else:
+        dist = (sorted_pred_goals[:, :, None, 0:2] - sorted_pred_goals[:, None, :, 0:2]).norm(dim=-1)
+        point_cover_mask = (dist < dist_thresh)
+
+    point_val = sorted_pred_scores.clone()  # (batch_size, N)
+    point_val_selected = torch.zeros_like(point_val)  # (batch_size, N)
+
+    ret_idxs = sorted_idxs.new_zeros(batch_size, num_ret_modes).long()
+    ret_trajs = sorted_pred_trajs.new_zeros(batch_size, num_ret_modes, num_timestamps, num_feat_dim)
+    ret_scores = sorted_pred_trajs.new_zeros(batch_size, num_ret_modes)
+    bs_idxs = torch.arange(batch_size).type_as(ret_idxs)
+
+    for k in range(num_ret_modes):
+        cur_idx = point_val.argmax(dim=-1)  # (batch_size)
+        ret_idxs[:, k] = cur_idx
+
+        new_cover_mask = point_cover_mask[bs_idxs, cur_idx]  # (batch_size, N)
+        point_val = point_val * (~new_cover_mask).float()  # (batch_size, N)
+        point_val_selected[bs_idxs, cur_idx] = -1
+        point_val += point_val_selected
+
+        ret_trajs[:, k] = sorted_pred_trajs[bs_idxs, cur_idx]
+        ret_scores[:, k] = sorted_pred_scores[bs_idxs, cur_idx]
+
+    bs_idxs = torch.arange(batch_size).type_as(sorted_idxs)[:, None].repeat(1, num_ret_modes)
+
+    ret_idxs = sorted_idxs[bs_idxs, ret_idxs]
+    return ret_trajs, ret_scores, ret_idxs
+
+
+def batch_nms_token(pred_trajs, pred_scores,
+                    dist_thresh, num_ret_modes=6,
+                    mode='static', speed=None):
+    """
+    Args:
+        pred_trajs (batch_size, num_modes, num_timestamps, 7)
+        pred_scores (batch_size, num_modes):
+        dist_thresh (float):
+        num_ret_modes (int, optional): Defaults to 6.
+
+    Returns:
+        ret_trajs (batch_size, num_ret_modes, num_timestamps, 5)
+        ret_scores (batch_size, num_ret_modes)
+        ret_idxs (batch_size, num_ret_modes)
+    """
+    batch_size, num_modes, num_feat_dim = pred_trajs.shape
+
+    sorted_idxs = pred_scores.argsort(dim=-1, descending=True)
+    bs_idxs_full = torch.arange(batch_size).type_as(sorted_idxs)[:, None].repeat(1, num_modes)
+    sorted_pred_scores = pred_scores[bs_idxs_full, sorted_idxs]
+    sorted_pred_goals = pred_trajs[bs_idxs_full, sorted_idxs]  # (batch_size, num_modes, num_timestamps, 7)
+
+    if mode == "nearby":
+        dist = (sorted_pred_goals[:, :, None, 0:2] - sorted_pred_goals[:, None, :, 0:2]).norm(dim=-1)
+        values, indices = torch.topk(dist, 5, dim=-1, largest=False)
+        thresh_hold = values[..., -1]
+        point_cover_mask = dist < thresh_hold[..., None]
+    else:
+        dist = (sorted_pred_goals[:, :, None, 0:2] - sorted_pred_goals[:, None, :, 0:2]).norm(dim=-1)
+        point_cover_mask = (dist < dist_thresh)
+
+    point_val = sorted_pred_scores.clone()  # (batch_size, N)
+    point_val_selected = torch.zeros_like(point_val)  # (batch_size, N)
+
+    ret_idxs = sorted_idxs.new_zeros(batch_size, num_ret_modes).long()
+    ret_goals = sorted_pred_goals.new_zeros(batch_size, num_ret_modes, num_feat_dim)
+    ret_scores = sorted_pred_goals.new_zeros(batch_size, num_ret_modes)
+    bs_idxs = torch.arange(batch_size).type_as(ret_idxs)
+
+    for k in range(num_ret_modes):
+        cur_idx = point_val.argmax(dim=-1)  # (batch_size)
+        ret_idxs[:, k] = cur_idx
+
+        new_cover_mask = point_cover_mask[bs_idxs, cur_idx]  # (batch_size, N)
+        point_val = point_val * (~new_cover_mask).float()  # (batch_size, N)
+        point_val_selected[bs_idxs, cur_idx] = -1
+        point_val += point_val_selected
+
+        ret_goals[:, k] = sorted_pred_goals[bs_idxs, cur_idx]
+        ret_scores[:, k] = sorted_pred_scores[bs_idxs, cur_idx]
+
+    bs_idxs = torch.arange(batch_size).type_as(sorted_idxs)[:, None].repeat(1, num_ret_modes)
+
+    ret_idxs = sorted_idxs[bs_idxs, ret_idxs]
+    return ret_goals, ret_scores, ret_idxs
+
+
+class TokenCls(Metric):
+
+    def __init__(self,
+                 max_guesses: int = 6,
+                 **kwargs) -> None:
+        super(TokenCls, self).__init__(**kwargs)
+        self.add_state('sum', default=torch.tensor(0.0), dist_reduce_fx='sum')
+        self.add_state('count', default=torch.tensor(0), dist_reduce_fx='sum')
+        self.max_guesses = max_guesses
+
+    def update(self,
+               pred: torch.Tensor,
+               target: torch.Tensor,
+               valid_mask: Optional[torch.Tensor] = None) -> None:
+        target = target[..., None]
+        acc = (pred[:, :self.max_guesses] == target).any(dim=1) * valid_mask
+        self.sum += acc.sum()
+        self.count += valid_mask.sum()
+
+    def compute(self) -> torch.Tensor:
+        return self.sum / self.count
+
+
+class minMultiFDE(Metric):
+
+    def __init__(self,
+                 max_guesses: int = 6,
+                 **kwargs) -> None:
+        super(minMultiFDE, self).__init__(**kwargs)
+        self.add_state('sum', default=torch.tensor(0.0), dist_reduce_fx='sum')
+        self.add_state('count', default=torch.tensor(0), dist_reduce_fx='sum')
+        self.max_guesses = max_guesses
+
+    def update(self,
+               pred: torch.Tensor,
+               target: torch.Tensor,
+               prob: Optional[torch.Tensor] = None,
+               valid_mask: Optional[torch.Tensor] = None,
+               keep_invalid_final_step: bool = True) -> None:
+        pred, target, prob, valid_mask, _ = valid_filter(pred, target, prob, valid_mask, None, keep_invalid_final_step)
+        pred_topk, _ = topk(self.max_guesses, pred, prob)
+        inds_last = (valid_mask * torch.arange(1, valid_mask.size(-1) + 1, device=self.device)).argmax(dim=-1)
+        self.sum += torch.norm(pred_topk[torch.arange(pred.size(0)), :, inds_last] -
+                               target[torch.arange(pred.size(0)), inds_last].unsqueeze(-2),
+                               p=2, dim=-1).min(dim=-1)[0].sum()
+        self.count += pred.size(0)
+
+    def compute(self) -> torch.Tensor:
+        return self.sum / self.count
+
+
+class minFDE(Metric):
+
+    def __init__(self,
+                 max_guesses: int = 6,
+                 **kwargs) -> None:
+        super(minFDE, self).__init__(**kwargs)
+        self.add_state('sum', default=torch.tensor(0.0), dist_reduce_fx='sum')
+        self.add_state('count', default=torch.tensor(0), dist_reduce_fx='sum')
+        self.max_guesses = max_guesses
+        self.eval_timestep = 70
+
+    def update(self,
+               pred: torch.Tensor,
+               target: torch.Tensor,
+               prob: Optional[torch.Tensor] = None,
+               valid_mask: Optional[torch.Tensor] = None,
+               keep_invalid_final_step: bool = True) -> None:
+        eval_timestep = min(self.eval_timestep, pred.shape[1]) - 1
+        self.sum += ((torch.norm(pred[:, eval_timestep-1:eval_timestep] - target[:, eval_timestep-1:eval_timestep], p=2, dim=-1) *
+                      valid_mask[:, eval_timestep-1].unsqueeze(1)).sum(dim=-1)).sum()
+        self.count += valid_mask[:, eval_timestep-1].sum()
+
+    def compute(self) -> torch.Tensor:
+        return self.sum / self.count
+
+
+class minMultiADE(Metric):
+
+    def __init__(self,
+                 max_guesses: int = 6,
+                 **kwargs) -> None:
+        super(minMultiADE, self).__init__(**kwargs)
+        self.add_state('sum', default=torch.tensor(0.0), dist_reduce_fx='sum')
+        self.add_state('count', default=torch.tensor(0), dist_reduce_fx='sum')
+        self.max_guesses = max_guesses
+
+    def update(self,
+               pred: torch.Tensor,
+               target: torch.Tensor,
+               prob: Optional[torch.Tensor] = None,
+               valid_mask: Optional[torch.Tensor] = None,
+               keep_invalid_final_step: bool = True,
+               min_criterion: str = 'FDE') -> None:
+        pred, target, prob, valid_mask, _ = valid_filter(pred, target, prob, valid_mask, None, keep_invalid_final_step)
+        pred_topk, _ = topk(self.max_guesses, pred, prob)
+        if min_criterion == 'FDE':
+            inds_last = (valid_mask * torch.arange(1, valid_mask.size(-1) + 1, device=self.device)).argmax(dim=-1)
+            inds_best = torch.norm(
+                pred_topk[torch.arange(pred.size(0)), :, inds_last] -
+                target[torch.arange(pred.size(0)), inds_last].unsqueeze(-2), p=2, dim=-1).argmin(dim=-1)
+            self.sum += ((torch.norm(pred_topk[torch.arange(pred.size(0)), inds_best] - target, p=2, dim=-1) *
+                          valid_mask).sum(dim=-1) / valid_mask.sum(dim=-1)).sum()
+        elif min_criterion == 'ADE':
+            self.sum += ((torch.norm(pred_topk - target.unsqueeze(1), p=2, dim=-1) *
+                          valid_mask.unsqueeze(1)).sum(dim=-1).min(dim=-1)[0] / valid_mask.sum(dim=-1)).sum()
+        else:
+            raise ValueError('{} is not a valid criterion'.format(min_criterion))
+        self.count += pred.size(0)
+
+    def compute(self) -> torch.Tensor:
+        return self.sum / self.count
+
+
+class minADE(Metric):
+
+    def __init__(self,
+                 max_guesses: int = 6,
+                 **kwargs) -> None:
+        super(minADE, self).__init__(**kwargs)
+        self.add_state('sum', default=torch.tensor(0.0), dist_reduce_fx='sum')
+        self.add_state('count', default=torch.tensor(0), dist_reduce_fx='sum')
+        self.max_guesses = max_guesses
+        self.eval_timestep = 70
+
+    def update(self,
+               pred: torch.Tensor,
+               target: torch.Tensor,
+               prob: Optional[torch.Tensor] = None,
+               valid_mask: Optional[torch.Tensor] = None,
+               keep_invalid_final_step: bool = True,
+               min_criterion: str = 'ADE') -> None:
+        # pred, target, prob, valid_mask, _ = valid_filter(pred, target, prob, valid_mask, None, keep_invalid_final_step)
+        # pred_topk, _ = topk(self.max_guesses, pred, prob)
+        # if min_criterion == 'FDE':
+        #     inds_last = (valid_mask * torch.arange(1, valid_mask.size(-1) + 1, device=self.device)).argmax(dim=-1)
+        #     inds_best = torch.norm(
+        #         pred[torch.arange(pred.size(0)), :, inds_last] -
+        #         target[torch.arange(pred.size(0)), inds_last].unsqueeze(-2), p=2, dim=-1).argmin(dim=-1)
+        #     self.sum += ((torch.norm(pred[torch.arange(pred.size(0)), inds_best] - target, p=2, dim=-1) *
+        #                   valid_mask).sum(dim=-1) / valid_mask.sum(dim=-1)).sum()
+        # elif min_criterion == 'ADE':
+        #     self.sum += ((torch.norm(pred - target.unsqueeze(1), p=2, dim=-1) *
+        #                   valid_mask.unsqueeze(1)).sum(dim=-1).min(dim=-1)[0] / valid_mask.sum(dim=-1)).sum()
+        # else:
+        #     raise ValueError('{} is not a valid criterion'.format(min_criterion))
+        eval_timestep = min(self.eval_timestep, pred.shape[1])
+        self.sum += ((torch.norm(pred[:, :eval_timestep] - target[:, :eval_timestep], p=2, dim=-1) * valid_mask[:, :eval_timestep]).sum(dim=-1) / pred.shape[1]).sum()
+        self.count += valid_mask[:, :eval_timestep].any(dim=-1).sum()
+
+    def compute(self) -> torch.Tensor:
+        return self.sum / self.count
+
+
+class AverageMeter(Metric):
+
+    def __init__(self, **kwargs) -> None:
+        super(AverageMeter, self).__init__(**kwargs)
+        self.add_state('sum', default=torch.tensor(0.0), dist_reduce_fx='sum')
+        self.add_state('count', default=torch.tensor(0), dist_reduce_fx='sum')
+
+    def update(self, val: torch.Tensor) -> None:
+        self.sum += val.sum()
+        self.count += val.numel()
+
+    def compute(self) -> torch.Tensor:
+        return self.sum / self.count
+
+
+class StateAccuracy(Metric):
+
+    def __init__(self, state_token: Dict[str, int], **kwargs) -> None:
+        super().__init__(**kwargs)
+        self.invalid_state = int(state_token['invalid'])
+        self.valid_state = int(state_token['valid'])
+        self.enter_state = int(state_token['enter'])
+        self.exit_state = int(state_token['exit'])
+
+        self.add_state('valid', default=torch.tensor(0), dist_reduce_fx='sum')
+        self.add_state('valid_count', default=torch.tensor(0), dist_reduce_fx='sum')
+        self.add_state('invalid', default=torch.tensor(0), dist_reduce_fx='sum')
+        self.add_state('invalid_count', default=torch.tensor(0), dist_reduce_fx='sum')
+
+    def update(self,
+               state_idx: torch.Tensor,
+               valid_mask: Optional[torch.Tensor] = None) -> None:
+
+        num_agent, num_step = state_idx.shape
+
+        # check the evaluation outputs
+        for a in range(num_agent):
+            bos_idx = torch.where(state_idx[a] == self.enter_state)[0]
+            eos_idx = torch.where(state_idx[a] == self.exit_state)[0]
+            bos = 0
+            eos = num_step - 1
+            if len(bos_idx) > 0:
+                bos = bos_idx[0]
+                self.invalid += (state_idx[a, :bos] == self.invalid_state).sum()
+                self.invalid_count += len(state_idx[a, :bos])
+            if len(eos_idx) > 0:
+                eos = eos_idx[0]
+                self.invalid += (state_idx[a, eos + 1:] == self.invalid_state).sum()
+                self.invalid_count += len(state_idx[a, eos + 1:])
+            self.valid += (state_idx[a, bos + 1 : eos] == self.valid_state).sum()
+            self.valid_count += len(state_idx[a, bos + 1 : eos])
+
+        # check the tokenization
+        if valid_mask is not None:
+
+            state_idx = state_idx.roll(shifts=1, dims=1)
+
+            for a in range(num_agent):
+                bos_idx = torch.where(state_idx[a] == self.enter_state)[0]
+                eos_idx = torch.where(state_idx[a] == self.exit_state)[0]
+                bos = 0
+                eos = num_step - 1
+                if len(bos_idx) > 0:
+                    bos = bos_idx[0]
+                    self.invalid += (valid_mask[a, :bos] == 0).sum()
+                    self.invalid_count += len(valid_mask[a, :bos])
+                if len(eos_idx) > 0:
+                    eos = eos_idx[-1]
+                    self.invalid += (valid_mask[a, eos + 1:] != 0).sum()
+                    self.invalid_count += len(valid_mask[a, eos + 1:])
+                self.invalid += (((state_idx[a, bos : eos + 1] > 0) != valid_mask[a, bos : eos + 1])[valid_mask[a, bos : eos + 1] == 0]).sum()
+                self.invalid_count += (valid_mask[a, bos : eos + 1] == 0).sum()
+                self.valid += (((state_idx[a, bos : eos + 1] > 0) != valid_mask[a, bos : eos + 1])[valid_mask[a, bos : eos + 1] == 1]).sum()
+                self.valid_count += (valid_mask[a, bos : eos + 1] == 1).sum()
+
+    def compute(self) -> Dict[str, torch.Tensor]:
+        return {'valid': self.valid / self.valid_count,
+                'invalid': self.invalid / self.invalid_count,
+                }
+
+    def __repr__(self):
+        head = "Results of " + self.__class__.__name__
+        results = self.compute()
+        body = [
+            "valid: {}".format(results['valid']),
+            "invalid: {}".format(results['invalid']),
+        ]
+        _repr_indent = 4
+        lines = [head] + [" " * _repr_indent + line for line in body]
+        return "\n".join(lines)
+
+
+class GridOverlapRate(Metric):
+
+    def __init__(self, num_step, state_token, seed_size, **kwargs) -> None:
+        super().__init__(**kwargs)
+        self.num_step = num_step
+        self.enter_state = int(state_token['enter'])
+        self.seed_size = seed_size
+        self.add_state('num_overlap_t', default=torch.zeros(num_step).long(), dist_reduce_fx='sum')
+        self.add_state('num_insert_agent_t', default=torch.zeros(num_step).long(), dist_reduce_fx='sum')
+        self.add_state('num_total_agent_t', default=torch.zeros(num_step).long(), dist_reduce_fx='sum')
+        self.add_state('num_exceed_seed_t', default=torch.zeros(num_step).long(), dist_reduce_fx='sum')
+
+    def update(self,
+               state_token: torch.Tensor,
+               grid_index: torch.Tensor) -> None:
+
+        for t in range(self.num_step):
+            inrange_mask_t = grid_index[:, t] != -1
+            insert_mask_t = (state_token[:, t] == self.enter_state) & inrange_mask_t
+            self.num_total_agent_t[t] += inrange_mask_t.sum()
+            self.num_insert_agent_t[t] += insert_mask_t.sum()
+            self.num_exceed_seed_t[t] += int(insert_mask_t.sum() >= self.seed_size)
+
+            occupied_grids = set(grid_index[:, t][(grid_index[:, t] != -1) & (state_token[:, t] != self.enter_state)].tolist())
+            to_inserted_grids = grid_index[:, t][(grid_index[:, t] != -1) & (state_token[:, t] == self.enter_state)].tolist()
+            while to_inserted_grids:
+                grid_index_t_i = to_inserted_grids.pop()
+                if grid_index_t_i in occupied_grids:
+                    self.num_overlap_t[t] += 1
+                occupied_grids.add(grid_index_t_i)
+
+    def compute(self) -> Dict[str, torch.Tensor]:
+        overlap_rate_t = self.num_overlap_t / self.num_insert_agent_t
+        overlap_rate_t.nan_to_num_()
+        return {'num_overlap_t': self.num_overlap_t,
+                'num_insert_agent_t': self.num_insert_agent_t,
+                'num_total_agent_t': self.num_total_agent_t,
+                'overlap_rate_t': overlap_rate_t,
+                'num_exceed_seed_t': self.num_exceed_seed_t,
+                }
+
+    def __repr__(self):
+        head = "Results of " + self.__class__.__name__
+        results = self.compute()
+        body = [
+            "num_overlap_t: {}".format(results['num_overlap_t'].tolist()),
+            "num_insert_agent_t: {}".format(results['num_insert_agent_t'].tolist()),
+            "num_total_agent_t: {}".format(results['num_total_agent_t'].tolist()),
+            "overlap_rate_t: {}".format(results['overlap_rate_t'].tolist()),
+            "num_exceed_seed_t: {}".format(results['num_exceed_seed_t'].tolist()),
+        ]
+        _repr_indent = 4
+        lines = [head] + [" " * _repr_indent + line for line in body]
+        return "\n".join(lines)
+
+
+class NumInsertAccuracy(Metric):
+
+    def __init__(self, state_token: Dict[str, int], **kwargs) -> None:
+        super().__init__(**kwargs)
+        self.invalid_state = int(state_token['invalid'])
+        self.valid_state = int(state_token['valid'])
+        self.enter_state = int(state_token['enter'])
+        self.exit_state = int(state_token['exit'])
+
+        self.add_state('valid', default=torch.tensor(0), dist_reduce_fx='sum')
+        self.add_state('valid_count', default=torch.tensor(0), dist_reduce_fx='sum')
+        self.add_state('invalid', default=torch.tensor(0), dist_reduce_fx='sum')
+        self.add_state('invalid_count', default=torch.tensor(0), dist_reduce_fx='sum')
+
+    def update(self,
+               state_idx: torch.Tensor,
+               valid_mask: Optional[torch.Tensor] = None) -> None:
+
+        num_agent, num_step = state_idx.shape
+
+        # check the evaluation outputs
+        for a in range(num_agent):
+            bos_idx = torch.where(state_idx[a] == self.enter_state)[0]
+            eos_idx = torch.where(state_idx[a] == self.exit_state)[0]
+            bos = 0
+            eos = num_step - 1
+            if len(bos_idx) > 0:
+                bos = bos_idx[0]
+                self.invalid += (state_idx[a, :bos] == self.invalid_state).sum()
+                self.invalid_count += len(state_idx[a, :bos])
+            if len(eos_idx) > 0:
+                eos = eos_idx[0]
+                self.invalid += (state_idx[a, eos + 1:] == self.invalid_state).sum()
+                self.invalid_count += len(state_idx[a, eos + 1:])
+            self.valid += (state_idx[a, bos + 1 : eos] == self.valid_state).sum()
+            self.valid_count += len(state_idx[a, bos + 1 : eos])
+
+        # check the tokenization
+        if valid_mask is not None:
+
+            state_idx = state_idx.roll(shifts=1, dims=1)
+
+            for a in range(num_agent):
+                bos_idx = torch.where(state_idx[a] == self.enter_state)[0]
+                eos_idx = torch.where(state_idx[a] == self.exit_state)[0]
+                bos = 0
+                eos = num_step - 1
+                if len(bos_idx) > 0:
+                    bos = bos_idx[0]
+                    self.invalid += (valid_mask[a, :bos] == 0).sum()
+                    self.invalid_count += len(valid_mask[a, :bos])
+                if len(eos_idx) > 0:
+                    eos = eos_idx[-1]
+                    self.invalid += (valid_mask[a, eos + 1:] != 0).sum()
+                    self.invalid_count += len(valid_mask[a, eos + 1:])
+                self.invalid += (((state_idx[a, bos : eos + 1] > 0) != valid_mask[a, bos : eos + 1])[valid_mask[a, bos : eos + 1] == 0]).sum()
+                self.invalid_count += (valid_mask[a, bos : eos + 1] == 0).sum()
+                self.valid += (((state_idx[a, bos : eos + 1] > 0) != valid_mask[a, bos : eos + 1])[valid_mask[a, bos : eos + 1] == 1]).sum()
+                self.valid_count += (valid_mask[a, bos : eos + 1] == 1).sum()
+
+    def compute(self) -> Dict[str, torch.Tensor]:
+        return {'valid': self.valid / self.valid_count,
+                'invalid': self.invalid / self.invalid_count,
+                }
+
+    def __repr__(self):
+        head = "Results of " + self.__class__.__name__
+        results = self.compute()
+        body = [
+            "valid: {}".format(results['valid']),
+            "invalid: {}".format(results['invalid']),
+        ]
+        _repr_indent = 4
+        lines = [head] + [" " * _repr_indent + line for line in body]
+        return "\n".join(lines)

seq_10_150_3_3_encode_occ_separate_offsets_bs8_128_no_seqindex_long_lastvalid_ablation_grid/backups/dev/utils/visualization.py

@@ -0,0 +1,1145 @@
+import math
+import os
+import torch
+import pickle
+import matplotlib.pyplot as plt
+import tensorflow as tf
+import numpy as np
+import numpy.typing as npt
+import fnmatch
+import seaborn as sns
+import matplotlib.axes as Axes
+import matplotlib.transforms as mtransforms
+from PIL import Image
+from functools import wraps
+from typing import Sequence, Union, Optional
+from tqdm import tqdm
+from typing import List, Literal
+from argparse import ArgumentParser
+from scipy.ndimage.filters import gaussian_filter
+from matplotlib.patches import FancyBboxPatch, Polygon, Rectangle, Circle
+from matplotlib.collections import LineCollection
+from torch_geometric.data import HeteroData, Dataset
+from waymo_open_dataset.protos import scenario_pb2
+
+from dev.utils.func import CONSOLE
+from dev.modules.attr_tokenizer import Attr_Tokenizer
+from dev.datasets.preprocess import TokenProcessor, cal_polygon_contour, AGENT_TYPE
+from dev.datasets.scalable_dataset import WaymoTargetBuilder
+
+
+__all__ = ['plot_occ_grid', 'plot_interact_edge', 'plot_map_edge', 'plot_insert_grid', 'plot_binary_map',
+           'plot_map_token', 'plot_prob_seed', 'plot_scenario', 'get_heatmap', 'draw_heatmap', 'plot_val', 'plot_tokenize']
+
+
+def safe_run(func):
+
+    @wraps(func)
+    def wrapper1(*args, **kwargs):
+        try:
+            return func(*args, **kwargs)
+        except Exception as e:
+            print(e)
+            return
+
+    @wraps(func)
+    def wrapper2(*args, **kwargs):
+        return func(*args, **kwargs)
+
+    if int(os.getenv('DEBUG', 0)):
+        return wrapper2
+    else:
+        return wrapper1
+
+
+@safe_run
+def plot_occ_grid(scenario_id, occ, gt_occ=None, save_path='', mode='agent', prefix=''):
+
+    def generate_box_edges(matrix, find_value=1):
+        y, x = np.where(matrix == find_value)
+        edges = []
+
+        for xi, yi in zip(x, y):
+            edges.append([(xi - 0.5, yi - 0.5), (xi + 0.5, yi - 0.5)])
+            edges.append([(xi + 0.5, yi - 0.5), (xi + 0.5, yi + 0.5)])
+            edges.append([(xi + 0.5, yi + 0.5), (xi - 0.5, yi + 0.5)])
+            edges.append([(xi - 0.5, yi + 0.5), (xi - 0.5, yi - 0.5)])
+    
+        return edges
+
+    os.makedirs(save_path, exist_ok=True)
+    n = int(math.sqrt(occ.shape[-1]))
+
+    plot_n = 3
+    plot_t = 5
+
+    occ_list = []
+    for i in range(plot_n):
+        for j in range(plot_t):
+            occ_list.append(occ[i, j].reshape(n, n))
+
+    occ_gt_list = []
+    if gt_occ is not None:
+        for i in range(plot_n):
+            for j in range(plot_t):
+                occ_gt_list.append(gt_occ[i, j].reshape(n, n))
+
+    row_labels = [f'n={n}' for n in range(plot_n)]
+    col_labels = [f't={t}' for t in range(plot_t)]
+
+    fig, axes = plt.subplots(plot_n, plot_t, figsize=(9, 6))
+    plt.subplots_adjust(wspace=0.1, hspace=0.1)
+
+    for i, ax in enumerate(axes.flat):
+        # NOTE: do not set vmin and vamx!
+        ax.imshow(occ_list[i], cmap='viridis', interpolation='nearest')
+        ax.axis('off')
+
+        if occ_gt_list:
+            gt_edges = generate_box_edges(occ_gt_list[i])
+            gts = LineCollection(gt_edges, colors='blue', linewidths=0.5)
+            ax.add_collection(gts)
+            insert_edges = generate_box_edges(occ_gt_list[i], find_value=-1)
+            inserts = LineCollection(insert_edges, colors='red', linewidths=0.5)
+            ax.add_collection(inserts)
+
+        ax.add_patch(plt.Rectangle((-0.5, -0.5), occ_list[i].shape[1], occ_list[i].shape[0],
+                                linewidth=2, edgecolor='black', facecolor='none'))
+
+    for i, ax in enumerate(axes[:, 0]):
+        ax.annotate(row_labels[i], xy=(-0.1, 0.5), xycoords="axes fraction",
+                    fontsize=12, ha="right", va="center", rotation=0)
+
+    for j, ax in enumerate(axes[0, :]):
+        ax.annotate(col_labels[j], xy=(0.5, 1.05), xycoords="axes fraction",
+                    fontsize=12, ha="center", va="bottom")
+
+    plt.savefig(os.path.join(save_path, f'{prefix}{scenario_id}_occ_{mode}.png'), dpi=500, bbox_inches='tight')
+    plt.close()
+
+
+@safe_run
+def plot_interact_edge(edge_index, scenario_ids, batch_sizes, num_seed, num_step, save_path='interact_edge_map',
+                        **kwargs):
+
+    num_batch = len(scenario_ids)
+    batches = torch.cat([
+        torch.arange(num_batch).repeat_interleave(repeats=batch_sizes, dim=0),
+        torch.arange(num_batch).repeat_interleave(repeats=num_seed, dim=0),
+    ], dim=0).repeat(num_step).numpy()
+
+    num_agent = batch_sizes.sum() + num_seed * num_batch
+    batch_sizes = torch.nn.functional.pad(batch_sizes, (1, 0), mode='constant', value=0)
+    ptr = torch.cumsum(batch_sizes, dim=0)
+    # assume difference scenarios and different timestep have the same number of seed agents
+    ptr_seed = torch.tensor(np.array([0] + [num_seed] * num_batch), device=ptr.device)
+
+    all_av_index = None
+    if 'av_index' in kwargs:
+        all_av_index = kwargs.pop('av_index').cpu() - ptr[:-1]
+
+    is_bos = np.zeros((batch_sizes.sum(), num_step)).astype(np.bool_)
+    if 'is_bos' in kwargs:
+        is_bos = kwargs.pop('is_bos').cpu().numpy()
+
+    src_index = torch.unique(edge_index[1])
+    for idx, src in enumerate(tqdm(src_index)):
+
+        src_batch = batches[src]
+
+        src_row = src % num_agent
+        if src_row // batch_sizes.sum() > 0:
+            seed_row = src_row % batch_sizes.sum() - ptr_seed[src_batch]
+            src_row = batch_sizes[src_batch + 1] + seed_row
+        else:
+            src_row = src_row - ptr[src_batch]
+
+        src_col = src // (num_agent)
+        src_mask = np.zeros((batch_sizes[src_batch + 1] + num_seed, num_step))
+        src_mask[src_row, src_col] = 1
+
+        tgt_mask = np.zeros((src_mask.shape[0], num_step))
+        tgt_index = edge_index[0, edge_index[1] == src]
+        for tgt in tgt_index:
+
+            tgt_batch = batches[tgt]
+
+            tgt_row = tgt % num_agent
+            if tgt_row // batch_sizes.sum() > 0:
+                seed_row = tgt_row % batch_sizes.sum() - ptr_seed[tgt_batch]
+                tgt_row = batch_sizes[tgt_batch + 1] + seed_row
+            else:
+                tgt_row = tgt_row - ptr[tgt_batch]
+
+            tgt_col = tgt // num_agent
+            tgt_mask[tgt_row, tgt_col] = 1
+            assert tgt_batch == src_batch
+
+        selected_step = tgt_mask.sum(axis=0) > 0
+        if selected_step.sum() > 1:
+            print(f"\nidx={idx}", src.item(), src_row.item(), src_col.item())
+            print(selected_step)
+            print(edge_index[:, edge_index[1] == src].tolist())
+
+        if all_av_index is not None:
+            kwargs['av_index'] = int(all_av_index[src_batch])
+
+        t = kwargs.get('t', src_col)
+        n = kwargs.get('n', 0)
+        is_bos_batch = is_bos[ptr[src_batch] : ptr[src_batch + 1]]
+        plot_binary_map(src_mask, tgt_mask, save_path, suffix=f'_{scenario_ids[src_batch]}_{t:02d}_{n:02d}_{idx:04d}',
+                        is_bos=is_bos_batch, **kwargs)
+
+
+@safe_run
+def plot_map_edge(edge_index, pos_a, data, save_path='map_edge_map'):
+
+    map_points = data['map_point']['position'][:, :2].cpu().numpy()
+    token_pos = data['pt_token']['position'][:, :2].cpu().numpy()
+    token_heading = data['pt_token']['orientation'].cpu().numpy()
+    num_pt = token_pos.shape[0]
+    
+    agent_index = torch.unique(edge_index[1])
+    for i in tqdm(agent_index):
+        xy = pos_a[i].cpu().numpy()
+        pt_index = edge_index[0, edge_index[1] == i].cpu().numpy()
+        pt_index = pt_index % num_pt
+
+        plt.subplots_adjust(left=0.3, right=0.7, top=0.7, bottom=0.3)
+        _, ax = plt.subplots()
+        ax.set_axis_off()
+
+        plot_map_token(ax, map_points, token_pos[pt_index], token_heading[pt_index], colors='blue')
+
+        ax.scatter(xy[0], xy[1], s=0.5, c='red', edgecolors='none')
+
+        os.makedirs(save_path, exist_ok=True)
+        plt.savefig(os.path.join(save_path, f'map_{i}.png'), dpi=600, bbox_inches='tight')
+        plt.close()
+
+
+def get_heatmap(x, y, prob, s=3, bins=1000):
+    heatmap, xedges, yedges = np.histogram2d(x, y, bins=bins, weights=prob, density=True)
+
+    heatmap = gaussian_filter(heatmap, sigma=s)
+
+    extent = [xedges[0], xedges[-1], yedges[0], yedges[-1]]
+    return heatmap.T, extent
+
+
+@safe_run
+def draw_heatmap(vector, vector_prob, gt_idx):
+    fig, ax = plt.subplots(figsize=(10, 10))
+    vector_prob = vector_prob.cpu().numpy()
+
+    for j in range(vector.shape[0]):
+        if j in gt_idx:
+            color = (0, 0, 1)
+        else:
+            grey_scale = max(0, 0.9 - vector_prob[j])
+            color = (0.9, grey_scale, grey_scale)
+
+        # if lane[j, k, -1] == 0: continue
+        x0, y0, x1, y1, = vector[j, :4]
+        ax.plot((x0, x1), (y0, y1), color=color, linewidth=2)
+
+    return plt
+
+
+@safe_run
+def plot_insert_grid(scenario_id, prob, grid, ego_pos, map, save_path='', prefix='', inference=False, indices=None,
+                     all_t_in_one=False):
+
+    """
+        prob: float array of shape (num_step, num_grid)
+        grid: float array of shape (num_grid, 2)
+    """
+
+    os.makedirs(save_path, exist_ok=True)
+
+    n = int(math.sqrt(prob.shape[1]))
+
+    # grid = grid[:, np.newaxis] + ego_pos[np.newaxis, ...]
+    for t in range(ego_pos.shape[0]):
+
+        plt.subplots_adjust(left=0.3, right=0.7, top=0.7, bottom=0.3)
+        _, ax = plt.subplots()
+
+        # plot probability
+        prob_t = prob[t].reshape(n, n)
+        plt.imshow(prob_t, cmap='viridis', interpolation='nearest')
+
+        if indices is not None:
+            indice = indices[t]
+
+            if isinstance(indice, (int, float, np.int_)):
+                indice = [indice]
+
+            for _indice in indice:
+                if _indice == -1: continue
+
+                row = _indice // n
+                col = _indice % n
+
+                rect = Rectangle((col - 0.5, row - 0.5), 1, 1, edgecolor='red', facecolor='none', lw=2)
+                ax.add_patch(rect)
+
+        ax.grid(False)
+        ax.set_aspect('equal', adjustable='box')
+
+        plt.title('Prob of Rel Position Grid')
+        plt.savefig(os.path.join(save_path, f'{prefix}{scenario_id}_heat_map_{t}.png'), dpi=300, bbox_inches='tight')
+        plt.close()
+
+        if all_t_in_one:
+            break
+
+
+@safe_run
+def plot_insert_grid(scenario_id, prob, indices=None, save_path='', prefix='', inference=False):
+
+    """
+        prob: float array of shape (num_seed, num_step, num_grid)
+        grid: float array of shape (num_grid, 2)
+    """
+
+    os.makedirs(save_path, exist_ok=True)
+
+    n = int(math.sqrt(prob.shape[-1]))
+
+    plot_n = 3
+    plot_t = 5
+
+    prob_list = []
+    for i in range(plot_n):
+        for j in range(plot_t):
+            prob_list.append(prob[i, j].reshape(n, n))
+
+    indice_list = []
+    if indices is not None:
+        for i in range(plot_n):
+            for j in range(plot_t):
+                indice_list.append(indices[i, j])
+
+    row_labels = [f'n={n}' for n in range(plot_n)]
+    col_labels = [f't={t}' for t in range(plot_t)]
+
+    fig, axes = plt.subplots(plot_n, plot_t, figsize=(9, 6))
+    fig.suptitle('Prob of Insert Position Grid')
+    plt.subplots_adjust(wspace=0.1, hspace=0.1)
+
+    for i, ax in enumerate(axes.flat):
+        ax.imshow(prob_list[i], cmap='viridis', interpolation='nearest')
+        ax.axis('off')
+
+        if indice_list:
+            row = indice_list[i] // n
+            col = indice_list[i] % n
+            rect = Rectangle((col - .5, row - .5), 1, 1, edgecolor='red', facecolor='none', lw=2)
+            ax.add_patch(rect)
+
+        ax.add_patch(plt.Rectangle((-0.5, -0.5), prob_list[i].shape[1], prob_list[i].shape[0],
+                                linewidth=2, edgecolor='black', facecolor='none'))
+
+    for i, ax in enumerate(axes[:, 0]):
+        ax.annotate(row_labels[i], xy=(-0.1, 0.5), xycoords="axes fraction",
+                    fontsize=12, ha="right", va="center", rotation=0)
+
+    for j, ax in enumerate(axes[0, :]):
+        ax.annotate(col_labels[j], xy=(0.5, 1.05), xycoords="axes fraction",
+                    fontsize=12, ha="center", va="bottom")
+
+    ax.grid(False)
+    ax.set_aspect('equal', adjustable='box')
+
+    plt.savefig(os.path.join(save_path, f'{prefix}{scenario_id}_insert_map.png'), dpi=500, bbox_inches='tight')
+    plt.close()
+
+
+@safe_run
+def plot_binary_map(src_mask, tgt_mask, save_path='', suffix='', av_index=None, is_bos=None, **kwargs):
+
+    from matplotlib.colors import ListedColormap
+    os.makedirs(save_path, exist_ok=True)
+
+    fig, axes = plt.subplots(1, 2, figsize=(10, 8))
+
+    title = []
+    if kwargs.get('t', None) is not None:
+        t = kwargs['t']
+        title.append(f't={t}')
+
+    if kwargs.get('n', None) is not None:
+        n = kwargs['n']
+        title.append(f'n={n}')
+
+    plt.title(' '.join(title))
+
+    cmap = ListedColormap(['white', 'green'])
+    axes[0].imshow(src_mask, cmap=cmap, interpolation='nearest')
+
+    cmap = ListedColormap(['white', 'orange'])
+    axes[1].imshow(tgt_mask, cmap=cmap, interpolation='nearest')
+
+    if av_index is not None:
+        rect = Rectangle((-0.5, av_index - 0.5), src_mask.shape[1], 1, edgecolor='red', facecolor='none', lw=2)
+        axes[0].add_patch(rect)
+        rect = Rectangle((-0.5, av_index - 0.5), tgt_mask.shape[1], 1, edgecolor='red', facecolor='none', lw=2)
+        axes[1].add_patch(rect)
+
+    if is_bos is not None:
+        rows, cols = np.where(is_bos)
+        for row, col in zip(rows, cols):
+            rect = Rectangle((col - 0.5, row - 0.5), 1, 1, edgecolor='blue', facecolor='none', lw=1)
+            axes[0].add_patch(rect)
+            rect = Rectangle((col - 0.5, row - 0.5), 1, 1, edgecolor='blue', facecolor='none', lw=1)
+            axes[1].add_patch(rect)
+
+    for ax in axes:
+        ax.set_xticks(range(src_mask.shape[1] + 1), minor=False)
+        ax.set_yticks(range(src_mask.shape[0] + 1), minor=False)
+        ax.grid(which='major', color='gray', linestyle='--', linewidth=0.5)
+
+    plt.savefig(os.path.join(save_path, f'map{suffix}.png'), dpi=300, bbox_inches='tight')
+    plt.close()
+
+
+@safe_run
+def plot_prob_seed(scenario_id, prob, save_path, prefix='', indices=None):
+
+    os.makedirs(save_path, exist_ok=True)
+
+    plt.figure(figsize=(8, 5))
+    plt.imshow(prob, cmap='viridis', aspect='auto')
+    plt.colorbar()
+
+    plt.title('Seed Probability')
+
+    if indices is not None:
+
+        for col in range(indices.shape[1]):
+            for row in indices[:, col]:
+
+                if row == -1: continue
+
+                rect = Rectangle((col - 0.5, row - 0.5), 1, 1, edgecolor='red', facecolor='none', lw=2)
+                plt.gca().add_patch(rect)
+
+    plt.tight_layout()
+    plt.savefig(os.path.join(save_path, f'{prefix}{scenario_id}_prob_seed.png'), dpi=300, bbox_inches='tight')
+    plt.close()
+
+
+@safe_run
+def plot_raw():
+    plt.figure(figsize=(30, 30))
+    plt.rcParams['axes.facecolor']='white'
+
+    data_path = '/u/xiuyu/work/dev4/data/waymo/scenario/training'
+    os.makedirs("data/vis/raw/0/", exist_ok=True)
+    file_list = os.listdir(data_path)
+
+    for cnt_file, file in enumerate(file_list):
+        file_path = os.path.join(data_path, file)
+        dataset = tf.data.TFRecordDataset(file_path, compression_type='')
+        for scenario_idx, data in enumerate(dataset):
+            scenario = scenario_pb2.Scenario()
+            scenario.ParseFromString(bytearray(data.numpy()))
+            tqdm.write(f"scenario id: {scenario.scenario_id}")
+
+            # draw maps
+            for i in range(len(scenario.map_features)):
+
+                # draw lanes
+                if str(scenario.map_features[i].lane) != '':
+                    line_x = [z.x for z in scenario.map_features[i].lane.polyline]
+                    line_y = [z.y for z in scenario.map_features[i].lane.polyline]
+                    plt.scatter(line_x, line_y, c='g', s=5)
+                    plt.text(line_x[0], line_y[0], str(scenario.map_features[i].id), fontdict={'family': 'serif', 'size': 20, 'color': 'green'})
+
+                # draw road_edge
+                if str(scenario.map_features[i].road_edge) != '':
+                    road_edge_x = [polyline.x for polyline in scenario.map_features[i].road_edge.polyline]
+                    road_edge_y = [polyline.y for polyline in scenario.map_features[i].road_edge.polyline]
+                    plt.scatter(road_edge_x, road_edge_y)
+                    plt.text(road_edge_x[0], road_edge_y[0], scenario.map_features[i].road_edge.type, fontdict={'family': 'serif', 'size': 20, 'color': 'black'})
+                    if scenario.map_features[i].road_edge.type == 2:
+                        plt.scatter(road_edge_x, road_edge_y, c='k')
+                    elif scenario.map_features[i].road_edge.type == 3:
+                        plt.scatter(road_edge_x, road_edge_y, c='purple')
+                        print(scenario.map_features[i].road_edge)
+                    else:
+                        plt.scatter(road_edge_x, road_edge_y, c='k')
+
+                # draw road_line
+                if str(scenario.map_features[i].road_line) != '':
+                    road_line_x = [j.x for j in scenario.map_features[i].road_line.polyline]
+                    road_line_y = [j.y for j in scenario.map_features[i].road_line.polyline]
+                    if scenario.map_features[i].road_line.type == 7:
+                        plt.plot(road_line_x, road_line_y, c='y')
+                    elif scenario.map_features[i].road_line.type == 8:
+                        plt.plot(road_line_x, road_line_y, c='y') 
+                    elif scenario.map_features[i].road_line.type == 6:
+                        plt.plot(road_line_x, road_line_y, c='y')
+                    elif scenario.map_features[i].road_line.type == 1:
+                        for i in range(int(len(road_line_x) / 7)):
+                            plt.plot(road_line_x[i * 7 : 5 + i * 7], road_line_y[i * 7 : 5 + i * 7], color='w')
+                    elif scenario.map_features[i].road_line.type == 2:
+                        plt.plot(road_line_x, road_line_y, c='w')
+                    else:
+                        plt.plot(road_line_x, road_line_y, c='w')
+            
+            # draw tracks
+            scenario_has_invalid_tracks = False
+            for i in range(len(scenario.tracks)):
+                traj_x = [center.center_x for center in scenario.tracks[i].states]
+                traj_y = [center.center_y for center in scenario.tracks[i].states]
+                head = [center.heading for center in scenario.tracks[i].states]
+                valid = [center.valid for center in scenario.tracks[i].states]
+                print(valid)
+                if i == scenario.sdc_track_index:
+                    plt.scatter(traj_x[0], traj_y[0], s=140, c='r', marker='s')
+                    plt.scatter([x for x, v in zip(traj_x, valid) if v],
+                                [y for y, v in zip(traj_y, valid) if v], s=14, c='r')
+                    plt.scatter([x for x, v in zip(traj_x, valid) if not v],
+                                [y for y, v in zip(traj_y, valid) if not v], s=14, c='m')
+                else:
+                    plt.scatter(traj_x[0], traj_y[0], s=140, c='k', marker='s')
+                    plt.scatter([x for x, v in zip(traj_x, valid) if v],
+                                [y for y, v in zip(traj_y, valid) if v], s=14, c='b')
+                    plt.scatter([x for x, v in zip(traj_x, valid) if not v],
+                                [y for y, v in zip(traj_y, valid) if not v], s=14, c='m')
+                if valid.count(False) > 0:
+                    scenario_has_invalid_tracks = True
+            if scenario_has_invalid_tracks:
+                plt.savefig(f"scenario_{scenario_idx}_{scenario.scenario_id}.png")
+                plt.clf()
+                breakpoint()
+        break
+
+
+colors = [
+    ('#1f77b4', '#1a5a8a'),  # blue
+    ('#2ca02c', '#217721'),  # green
+    ('#ff7f0e', '#cc660b'),  # orange
+    ('#9467bd', '#6f4a91'),  # purple
+    ('#d62728', '#a31d1d'),  # red
+    ('#000000', '#000000'),  # black
+]
+
+@safe_run
+def plot_gif():
+    data_path = "/u/xiuyu/work/dev4/data/waymo_processed/training"
+    os.makedirs("data/vis/processed/0/gif", exist_ok=True)
+    file_list = os.listdir(data_path)
+
+    for scenario_idx, file in tqdm(enumerate(file_list), leave=False, desc="Scenario"):
+
+        fig, ax = plt.subplots()
+        ax.set_axis_off()
+
+        file_path = os.path.join(data_path, file)
+        data = pickle.load(open(file_path, "rb"))
+        scenario_id = data['scenario_id']
+
+        save_path = os.path.join("data/vis/processed/0/gif",
+                                 f"scenario_{scenario_idx}_{scenario_id}.gif")
+        if os.path.exists(save_path):
+            tqdm.write(f"Skipped {save_path}.")
+            continue
+
+        # draw maps
+        ax.scatter(data['map_point']['position'][:, 0],
+                   data['map_point']['position'][:, 1], s=0.2, c='black', edgecolors='none')
+
+        # draw agents
+        agent_data = data['agent']
+        av_index = agent_data['av_index']
+        position = agent_data['position'] # (num_agent, 91, 3)
+        heading = agent_data['heading'] # (num_agent, 91)
+        shape = agent_data['shape'] # (num_agent, 91, 3)
+        category = agent_data['category'] # (num_agent,)
+        valid_mask = (position[..., 0] != 0) & (position[..., 1] != 0) # (num_agent, 91)
+
+        num_agent = valid_mask.shape[0]
+        num_timestep = position.shape[1]
+        is_av = np.arange(num_agent) == int(av_index)
+
+        is_blue = valid_mask.sum(axis=1) == num_timestep
+        is_green = ~valid_mask[:, 0] & valid_mask[:, -1]
+        is_orange = valid_mask[:, 0] & ~valid_mask[:, -1]
+        is_purple = (valid_mask.sum(axis=1) != num_timestep
+                    ) & (~is_green) & (~is_orange)
+        agent_colors = np.zeros((num_agent,))
+        agent_colors[is_blue] = 1
+        agent_colors[is_green] = 2
+        agent_colors[is_orange] = 3
+        agent_colors[is_purple] = 4
+        agent_colors[is_av] = 5
+
+        veh_mask = category == 1
+        ped_mask = category == 2
+        cyc_mask = category == 3
+        shape[veh_mask, :, 1] = 1.8
+        shape[veh_mask, :, 0] = 1.8
+        shape[ped_mask, :, 1] = 0.5
+        shape[ped_mask, :, 0] = 0.5
+        shape[cyc_mask, :, 1] = 1.0
+        shape[cyc_mask, :, 0] = 1.0
+
+        fig_paths = []
+        for tid in tqdm(range(num_timestep), leave=False, desc="Timestep"):
+            current_valid_mask = valid_mask[:, tid]
+            xs = position[current_valid_mask, tid, 0]
+            ys = position[current_valid_mask, tid, 1]
+            widths = shape[current_valid_mask, tid, 1]
+            lengths = shape[current_valid_mask, tid, 0]
+            angles = heading[current_valid_mask, tid]
+            current_agent_colors = agent_colors[current_valid_mask]
+
+            drawn_agents = []
+            contours = cal_polygon_contour(xs, ys, angles, widths, lengths) # (num_agent, 4, 2)
+            contours = np.concatenate([contours, contours[:, 0:1]], axis=1) # (num_agent, 5, 2)
+            for x, y, width, length, angle, color_type in zip(
+                xs, ys, widths, lengths, angles, current_agent_colors):
+                agent = plt.Rectangle((x, y), width, length, angle=((angle + np.pi / 2) / np.pi * 360) % 360,
+                                      linewidth=0.2,
+                                      facecolor=colors[int(color_type) - 1][0],
+                                      edgecolor=colors[int(color_type) - 1][1])
+                ax.add_patch(agent)
+                drawn_agents.append(agent)
+            plt.gca().set_aspect('equal', adjustable='box')
+            # for contour, color_type in zip(contours, agent_colors):
+            #     drawn_agent = ax.plot(contour[:, 0], contour[:, 1])
+            #     drawn_agents.append(drawn_agent)
+
+            fig_path = os.path.join("data/vis/processed/0/",
+                                    f"scenario_{scenario_idx}_{scenario_id}_{tid}.png")
+            plt.savefig(fig_path, dpi=600)
+            fig_paths.append(fig_path)
+
+            for drawn_agent in drawn_agents:
+                drawn_agent.remove()
+
+        plt.close()
+
+        # generate gif
+        import imageio.v2 as imageio
+        images = []
+        for fig_path in tqdm(fig_paths, leave=False, desc="Generate gif ..."):
+            images.append(imageio.imread(fig_path))
+        imageio.mimsave(save_path, images, duration=0.1)
+
+
+@safe_run
+def plot_map_token(ax: Axes, map_points: npt.NDArray, token_pos: npt.NDArray, token_heading: npt.NDArray, colors: Union[str, npt.NDArray]=None):
+
+    plot_map(ax, map_points)
+
+    x, y = token_pos[:, 0], token_pos[:, 1]
+    u = np.cos(token_heading)
+    v = np.sin(token_heading)
+
+    if colors is None:
+        colors = np.random.rand(x.shape[0], 3)
+    ax.quiver(x, y, u, v, angles='xy', scale_units='xy', scale=0.2, color=colors, width=0.005,
+               headwidth=0.2, headlength=2)
+    ax.scatter(x, y, color='blue', s=0.2, edgecolors='none')
+    ax.axis("equal")
+
+
+@safe_run
+def plot_map(ax: Axes, map_points: npt.NDArray, color='black'):
+    ax.scatter(map_points[:, 0], map_points[:, 1], s=0.2, c=color, edgecolors='none')
+
+    xmin = np.min(map_points[:, 0])
+    xmax = np.max(map_points[:, 0])
+    ymin = np.min(map_points[:, 1])
+    ymax = np.max(map_points[:, 1])
+    ax.set_xlim(xmin, xmax)
+    ax.set_ylim(ymin, ymax)
+
+
+@safe_run
+def plot_agent(ax: Axes, xy: Sequence[float], heading: float, type: str, state, is_av: bool=False,
+               pl2seed_radius: float=25., attr_tokenizer: Attr_Tokenizer=None, enter_index: list=[], **kwargs):
+
+    if type == 'veh':
+        length = 4.3
+        width = 1.8
+        size = 1.0
+    elif type == 'ped':
+        length = 0.5
+        width = 0.5
+        size = 0.1
+    elif type == 'cyc':
+        length = 1.9
+        width = 0.5
+        size = 0.3
+    else:
+        raise ValueError(f"Unsupported agent type {type}")
+
+    if kwargs.get('label', None) is not None:
+        ax.text(
+            xy[0] + 1.5, xy[1] + 1.5,
+            kwargs['label'], fontsize=2, color="darkred", ha="center", va="center"
+        )
+
+    patch = FancyBboxPatch([-length / 2, -width / 2], length, width, linewidth=.2, **kwargs)
+    transform = (
+        mtransforms.Affine2D().rotate(heading).translate(xy[0], xy[1])
+        + ax.transData
+    )
+    patch.set_transform(transform)
+
+    kwargs['label'] = None
+    angles = [0, 2 * np.pi / 3, np.pi, 4 * np.pi / 3]
+    pts = np.stack([size * np.cos(angles), size * np.sin(angles)], axis=-1)
+    center_patch = Polygon(pts, zorder=10., linewidth=.2, **kwargs)
+    center_patch.set_transform(transform)
+
+    ax.add_patch(patch)
+    ax.add_patch(center_patch)
+
+    if is_av:
+
+        if attr_tokenizer is not None:
+
+            circle_patch = Circle(
+                (xy[0], xy[1]), pl2seed_radius, linewidth=0.5, edgecolor='gray', linestyle='--', facecolor='none'
+            )
+            ax.add_patch(circle_patch)
+
+            grid = attr_tokenizer.get_grid(torch.tensor(np.array(xy)).float(),
+                                           torch.tensor(np.array([heading])).float()).numpy()[0] # (num_grid, 2)
+            ax.scatter(grid[:, 0], grid[:, 1], s=0.3, c='blue', edgecolors='none')
+            ax.text(grid[0, 0], grid[0, 1], 'Front', fontsize=2, color='darkred', ha='center', va='center')
+            ax.text(grid[-1, 0], grid[-1, 1], 'Back', fontsize=2, color='darkred', ha='center', va='center')
+
+        if enter_index:
+            for i in enter_index:
+                ax.plot(grid[int(i), 0], grid[int(i), 1], marker='x', color='red', markersize=1)
+
+    return patch, center_patch
+
+
+@safe_run
+def plot_all(map, xs, ys, angles, types, colors, is_avs, pl2seed_radius: float=25.,
+             attr_tokenizer: Attr_Tokenizer=None, enter_index: list=[], labels: list=[], **kwargs):
+
+    plt.subplots_adjust(left=0.3, right=0.7, top=0.7, bottom=0.3)
+    _, ax = plt.subplots()
+    ax.set_axis_off()
+
+    plot_map(ax, map)
+
+    if not labels:
+        labels = [None] * xs.shape[0]
+
+    for x, y, angle, type, color, label, is_av in zip(xs, ys, angles, types, colors, labels, is_avs):
+        assert type in ('veh', 'ped', 'cyc'), f"Unsupported type {type}."
+        plot_agent(ax, [x, y], angle.item(), type, None, is_av, facecolor=color, edgecolor='k', label=label,
+                   pl2seed_radius=pl2seed_radius, attr_tokenizer=attr_tokenizer, enter_index=enter_index)
+
+    ax.grid(False)
+    ax.set_aspect('equal', adjustable='box')
+
+    # ax.legend(loc='best', frameon=True)
+
+    if kwargs.get('save_path', None):
+        plt.savefig(kwargs['save_path'], dpi=600, bbox_inches="tight")
+
+    plt.close()
+
+    return ax
+
+
+@safe_run
+def plot_file(gt_folder: str,
+              folder: Optional[str] = None,
+              files: Optional[str] = None):
+    from dev.metrics.compute_metrics import _unbatch
+
+    if files is None:
+        assert os.path.exists(folder), f'Path {folder} does not exist.'
+        files = list(fnmatch.filter(os.listdir(folder), 'idx_*_rollouts.pkl'))
+        CONSOLE.log(f'Found {len(files)} rollouts files from {folder}.')
+
+
+    if folder is None:
+        assert os.path.exists(files), f'Path {files} does not exist.'
+        folder = os.path.dirname(files)
+        files = [files]
+
+    parent, folder_name = os.path.split(folder.rstrip(os.sep))
+    save_path = os.path.join(parent, f'{folder_name}_plots')
+
+    for file in (pbar := tqdm(files, leave=False, desc='Plotting files ...')):
+        pbar.set_postfix(file=file)
+
+        with open(os.path.join(folder, file), 'rb') as f:
+            preds = pickle.load(f)
+
+        scenario_ids = preds['_scenario_id']
+        agent_batch = preds['agent_batch']
+        agent_id = _unbatch(preds['agent_id'], agent_batch)
+        preds_traj = _unbatch(preds['pred_traj'], agent_batch)
+        preds_head = _unbatch(preds['pred_head'], agent_batch)
+        preds_type = _unbatch(preds['pred_type'], agent_batch)
+        preds_state = _unbatch(preds['pred_state'], agent_batch)
+        preds_valid = _unbatch(preds['pred_valid'], agent_batch)
+
+        for i, scenario_id in enumerate(scenario_ids):
+            n_rollouts = preds_traj[0].shape[1]
+
+            for j in range(n_rollouts):  # 1
+                pred = dict(scenario_id=[scenario_id],
+                            pred_traj=preds_traj[i][:, j],
+                            pred_head=preds_head[i][:, j],
+                            pred_state=preds_state[i][:, j],
+                            pred_type=preds_type[i][:, j],
+                    )
+                av_index = agent_id[i][:, 0].tolist().index(preds['av_id'])  # NOTE: hard code!!!
+
+                data_path = os.path.join(gt_folder, 'validation', f'{scenario_id}.pkl')
+                with open(data_path, 'rb') as f:
+                    data = pickle.load(f)
+                plot_val(data, pred, av_index=av_index, save_path=save_path)
+
+
+@safe_run
+def plot_val(data: Union[dict, str], pred: dict, av_index: int, save_path: str, suffix: str='',
+             pl2seed_radius: float=75., attr_tokenizer=None, **kwargs):
+
+    if isinstance(data, str):
+        assert data.endswith('.pkl'), f'Got invalid data path {data}.'
+        assert os.path.exists(data), f'Path {data} does not exist.'
+        with open(data, 'rb') as f:
+            data = pickle.load(f)
+
+    map_point = data['map_point']['position'].cpu().numpy()
+
+    scenario_id = pred['scenario_id'][0]
+    pred_traj = pred['pred_traj'].cpu().numpy() # (num_agent, num_future_step, 2)
+    pred_type = list(map(lambda i: AGENT_TYPE[i], pred['pred_type'].tolist()))
+    pred_state = pred['pred_state'].cpu().numpy()
+    pred_head = pred['pred_head'].cpu().numpy()
+    ids = np.arange(pred_traj.shape[0])
+
+    if 'agent_labels' in pred:
+        kwargs.update(agent_labels=pred['agent_labels'])
+
+    plot_scenario(scenario_id, map_point, pred_traj, pred_head, pred_state, pred_type,
+                  av_index=av_index, ids=ids, save_path=save_path, suffix=suffix,
+                  pl2seed_radius=pl2seed_radius, attr_tokenizer=attr_tokenizer, **kwargs)
+
+
+@safe_run
+def plot_scenario(scenario_id: str,
+                  map_data: npt.NDArray,
+                  traj: npt.NDArray,
+                  heading: npt.NDArray,
+                  state: npt.NDArray,
+                  types: List[str],
+                  av_index: int,
+                  color_type: Literal['state', 'type', 'seed']='seed',
+                  state_type: List[str]=['invalid', 'valid', 'enter', 'exit'],
+                  plot_enter: bool=False,
+                  suffix: str='',
+                  pl2seed_radius: float=25.,
+                  attr_tokenizer: Attr_Tokenizer=None,
+                  enter_index: List[list] = [],
+                  save_gif: bool=True,
+                  tokenized: bool=False,
+                  agent_labels: List[List[Optional[str]]] = [],
+                  **kwargs):
+
+    num_historical_steps = 11
+    shift = 5
+    num_agent, num_timestep = traj.shape[:2]
+
+    if tokenized:
+        num_historical_steps = 2
+        shift = 1
+
+    if 'save_path' in kwargs and kwargs['save_path'] != '':
+        os.makedirs(kwargs['save_path'], exist_ok=True)
+        save_id = int(max([0] + list(map(lambda fname: int(fname.split("_")[-1]),
+                                        filter(lambda fname: fname.startswith(scenario_id)
+                                               and os.path.isdir(os.path.join(kwargs['save_path'], fname)),
+                                        os.listdir(kwargs['save_path'])))))) + 1
+        os.makedirs(f"{kwargs['save_path']}/{scenario_id}_{str(save_id).zfill(3)}", exist_ok=True)
+
+        if save_id > 1:
+            try:
+                import shutil
+                shutil.rmtree(f"{kwargs['save_path']}/{scenario_id}_{str(save_id - 1).zfill(3)}")
+            except:
+                pass
+
+    visible_mask = state != state_type.index('invalid')
+    if not plot_enter:
+        visible_mask &= (state != state_type.index('enter'))
+
+    last_valid_step = visible_mask.shape[1] - 1 - torch.argmax(torch.Tensor(visible_mask).flip(dims=[1]).long(), dim=1)
+    ids = None
+    if 'ids' in kwargs:
+        ids = kwargs['ids']
+        last_valid_step = {int(ids[i]): int(last_valid_step[i]) for i in range(len(ids))}
+
+    # agent colors
+    agent_colors = np.zeros((num_agent, num_timestep, 3))
+
+    agent_palette = sns.color_palette('husl', n_colors=7)
+    state_colors = {state: np.array(agent_palette[i]) for i, state in enumerate(state_type)}
+    seed_colors = {seed: np.array(agent_palette[i]) for i, seed in enumerate(['existing', 'entered', 'exited'])}
+
+    if color_type == 'state':
+        for t in range(state.shape[1]):
+            agent_colors[state[:, t] == state_type.index('invalid'), t * shift : (t + 1) * shift] = state_colors['invalid']
+            agent_colors[state[:, t] == state_type.index('valid'), t * shift : (t + 1) * shift] = state_colors['valid']
+            agent_colors[state[:, t] == state_type.index('enter'), t * shift : (t + 1) * shift] = state_colors['enter']
+            agent_colors[state[:, t] == state_type.index('exit'), t * shift : (t + 1) * shift] = state_colors['exit']
+
+    if color_type == 'seed':
+        agent_colors[:, :] = seed_colors['existing']
+        is_exited = np.any(state[:, num_historical_steps - 1:] == state_type.index('exit'), axis=-1)
+        is_entered = np.any(state[:, num_historical_steps - 1:] == state_type.index('enter'), axis=-1)
+        is_entered[av_index + 1:] = True  # NOTE: hard code, need improvment
+        agent_colors[is_exited, :] = seed_colors['exited']
+        agent_colors[is_entered, :] = seed_colors['entered']
+
+    agent_colors[av_index, :] = np.array(agent_palette[-1])
+    is_av = np.zeros_like(state[:, 0]).astype(np.bool_)
+    is_av[av_index] = True
+
+    # draw agents
+    fig_paths = []
+    for tid in tqdm(range(num_timestep), leave=False, desc="Plot ..."):
+        mask_t = visible_mask[:, tid]
+        xs = traj[mask_t, tid, 0]
+        ys = traj[mask_t, tid, 1]
+        angles = heading[mask_t, tid]
+        colors = agent_colors[mask_t, tid]
+        types_t = [types[i] for i, mask in enumerate(mask_t) if mask]
+        if ids is not None:
+            ids_t = ids[mask_t]
+        is_av_t = is_av[mask_t]
+        enter_index_t = enter_index[tid] if enter_index else None
+        labels = []
+        if agent_labels:
+            labels = [agent_labels[i][tid // shift] for i in range(len(agent_labels)) if mask_t[i]]
+
+        fig_path = None
+        if 'save_path' in kwargs:
+            save_path = kwargs['save_path']
+            fig_path = os.path.join(f"{save_path}/{scenario_id}_{str(save_id).zfill(3)}", f"{tid}.png")
+            fig_paths.append(fig_path)
+
+        plot_all(map_data, xs, ys, angles, types_t, colors=colors, save_path=fig_path, is_avs=is_av_t,
+                 pl2seed_radius=pl2seed_radius, attr_tokenizer=attr_tokenizer, enter_index=enter_index_t, labels=labels)
+
+    # generate gif
+    if fig_paths and save_gif:
+        os.makedirs(os.path.join(save_path, 'gifs'), exist_ok=True)
+        images = []
+        gif_path = f"{save_path}/gifs/{scenario_id}_{str(save_id).zfill(3)}.gif"
+        for fig_path in tqdm(fig_paths, leave=False, desc="Generate gif ..."):
+            images.append(Image.open(fig_path))
+        try:
+            images[0].save(gif_path, save_all=True, append_images=images[1:], duration=100, loop=0)
+            tqdm.write(f"Saved gif at {gif_path}")
+            try:
+                import shutil
+                shutil.rmtree(f"{save_path}/{scenario_id}_{str(save_id).zfill(3)}")
+                os.remove(f"{save_path}/gifs/{scenario_id}_{str(save_id - 1).zfill(3)}.gif")
+            except:
+                pass
+        except Exception as e:
+            tqdm.write(f"{e}! Failed to save gif at {gif_path}")
+
+
+def match_token_map(data):
+
+    # init map token
+    argmin_sample_len = 3
+    map_token_traj_path = '/u/xiuyu/work/dev4/dev/tokens/map_traj_token5.pkl'
+
+    map_token_traj = pickle.load(open(map_token_traj_path, 'rb'))
+    map_token = {'traj_src': map_token_traj['traj_src'], }
+    traj_end_theta = np.arctan2(map_token['traj_src'][:, -1, 1] - map_token['traj_src'][:, -2, 1],
+                                map_token['traj_src'][:, -1, 0] - map_token['traj_src'][:, -2, 0])
+    indices = torch.linspace(0, map_token['traj_src'].shape[1]-1, steps=argmin_sample_len).long()
+    map_token['sample_pt'] = torch.from_numpy(map_token['traj_src'][:, indices]).to(torch.float)
+    map_token['traj_end_theta'] = torch.from_numpy(traj_end_theta).to(torch.float)
+    map_token['traj_src'] = torch.from_numpy(map_token['traj_src']).to(torch.float)
+
+    traj_pos = data['map_save']['traj_pos'].to(torch.float)
+    traj_theta = data['map_save']['traj_theta'].to(torch.float)
+    pl_idx_list = data['map_save']['pl_idx_list']
+    token_sample_pt = map_token['sample_pt'].to(traj_pos.device)
+    token_src = map_token['traj_src'].to(traj_pos.device)
+    max_traj_len = map_token['traj_src'].shape[1]
+    pl_num = traj_pos.shape[0]
+
+    pt_token_pos = traj_pos[:, 0, :].clone()
+    pt_token_orientation = traj_theta.clone()
+    cos, sin = traj_theta.cos(), traj_theta.sin()
+    rot_mat = traj_theta.new_zeros(pl_num, 2, 2)
+    rot_mat[..., 0, 0] = cos
+    rot_mat[..., 0, 1] = -sin
+    rot_mat[..., 1, 0] = sin
+    rot_mat[..., 1, 1] = cos
+    traj_pos_local = torch.bmm((traj_pos - traj_pos[:, 0:1]), rot_mat.view(-1, 2, 2))
+    distance = torch.sum((token_sample_pt[None] - traj_pos_local.unsqueeze(1)) ** 2, dim=(-2, -1))
+    pt_token_id = torch.argmin(distance, dim=1)
+
+    noise = False
+    if noise:
+        topk_indices = torch.argsort(torch.sum((token_sample_pt[None] - traj_pos_local.unsqueeze(1)) ** 2, dim=(-2, -1)), dim=1)[:, :8]
+        sample_topk = torch.randint(0, topk_indices.shape[-1], size=(topk_indices.shape[0], 1), device=topk_indices.device)
+        pt_token_id = torch.gather(topk_indices, 1, sample_topk).squeeze(-1)
+
+    # cos, sin = traj_theta.cos(), traj_theta.sin()
+    # rot_mat = traj_theta.new_zeros(pl_num, 2, 2)
+    # rot_mat[..., 0, 0] = cos
+    # rot_mat[..., 0, 1] = sin
+    # rot_mat[..., 1, 0] = -sin
+    # rot_mat[..., 1, 1] = cos
+    # token_src_world = torch.bmm(token_src[None, ...].repeat(pl_num, 1, 1, 1).reshape(pl_num, -1, 2),
+    #                             rot_mat.view(-1, 2, 2)).reshape(pl_num, token_src.shape[0], max_traj_len, 2) + traj_pos[:, None, [0], :]
+    # token_src_world_select = token_src_world.view(-1, 1024, 11, 2)[torch.arange(pt_token_id.view(-1).shape[0]), pt_token_id.view(-1)].view(pl_num, max_traj_len, 2)
+
+    pl_idx_full = pl_idx_list.clone()
+    token2pl = torch.stack([torch.arange(len(pl_idx_list), device=traj_pos.device), pl_idx_full.long()])
+    count_nums = []
+    for pl in pl_idx_full.unique():
+        pt = token2pl[0, token2pl[1, :] == pl]
+        left_side = (data['pt_token']['side'][pt] == 0).sum()
+        right_side = (data['pt_token']['side'][pt] == 1).sum()
+        center_side = (data['pt_token']['side'][pt] == 2).sum()
+        count_nums.append(torch.Tensor([left_side, right_side, center_side]))
+    count_nums = torch.stack(count_nums, dim=0)
+    num_polyline = int(count_nums.max().item())
+    traj_mask = torch.zeros((int(len(pl_idx_full.unique())), 3, num_polyline), dtype=bool)
+    idx_matrix = torch.arange(traj_mask.size(2)).unsqueeze(0).unsqueeze(0)
+    idx_matrix = idx_matrix.expand(traj_mask.size(0), traj_mask.size(1), -1)
+    counts_num_expanded = count_nums.unsqueeze(-1)
+    mask_update = idx_matrix < counts_num_expanded
+    traj_mask[mask_update] = True
+
+    data['pt_token']['traj_mask'] = traj_mask
+    data['pt_token']['position'] = torch.cat([pt_token_pos, torch.zeros((data['pt_token']['num_nodes'], 1),
+                                                                        device=traj_pos.device, dtype=torch.float)], dim=-1)
+    data['pt_token']['orientation'] = pt_token_orientation
+    data['pt_token']['height'] = data['pt_token']['position'][:, -1]
+    data[('pt_token', 'to', 'map_polygon')] = {}
+    data[('pt_token', 'to', 'map_polygon')]['edge_index'] = token2pl  # (2, num_points)
+    data['pt_token']['token_idx'] = pt_token_id
+    return data
+
+
+@safe_run
+def plot_tokenize(data, save_path: str):
+
+    shift = 5
+    token_size = 2048
+    pl2seed_radius = 75
+
+    # transformation
+    transform = WaymoTargetBuilder(num_historical_steps=11,
+                                   num_future_steps=80,
+                                   max_num=32,
+                                   training=False)
+
+    grid_range = 150.
+    grid_interval = 3.
+    angle_interval = 3.
+    attr_tokenizer = Attr_Tokenizer(grid_range=grid_range,
+                                    grid_interval=grid_interval,
+                                    radius=pl2seed_radius,
+                                    angle_interval=angle_interval)
+
+    # tokenization
+    token_processor = TokenProcessor(token_size,
+                                     training=False,
+                                     predict_motion=True,
+                                     predict_state=True,
+                                     predict_map=True,
+                                     state_token={'invalid': 0, 'valid': 1, 'enter': 2, 'exit': 3},
+                                     pl2seed_radius=pl2seed_radius)
+    CONSOLE.log(f"Loaded token processor with token_size: {token_size}")
+
+    # preprocess
+    data: HeteroData = transform(data)
+    tokenized_data = token_processor(data)
+    CONSOLE.log(f"Keys in tokenized data:\n{tokenized_data.keys()}")
+
+    # plot
+    agent_data = tokenized_data['agent']
+    map_data = tokenized_data['map_point']
+    # CONSOLE.log(f"Keys in agent data:\n{agent_data.keys()}")
+
+    av_index = agent_data['av_index']
+    raw_traj = agent_data['position'][..., :2].contiguous()  # [n_agent, n_step, 2]
+    raw_heading = agent_data['heading']  # [n_agent, n_step]
+
+    traj = agent_data['traj_pos'][..., :2].contiguous()  # [n_agent, n_step, 6, 2]
+    traj = traj[:, :, 1:, :].flatten(1, 2)
+    traj = torch.cat([raw_traj[:, :1], traj], dim=1)
+    heading = agent_data['traj_heading']  # [n_agent, n_step, 6]
+    heading = heading[:, :, 1:].flatten(1, 2)
+    heading = torch.cat([raw_heading[:, :1], heading], dim=1)
+
+    agent_state = agent_data['state_idx'].repeat_interleave(repeats=shift, dim=-1)
+    agent_state = torch.cat([torch.zeros_like(agent_state[:, :1]), agent_state], dim=1)
+    agent_type = agent_data['type']
+    ids = np.arange(raw_traj.shape[0])
+
+    plot_scenario(scenario_id=tokenized_data['scenario_id'],
+                  map_data=tokenized_data['map_point']['position'].numpy(),
+                  traj=raw_traj.numpy(),
+                  heading=raw_heading.numpy(),
+                  state=agent_state.numpy(),
+                  types=list(map(lambda i: AGENT_TYPE[i], agent_type.tolist())),
+                  av_index=av_index,
+                  ids=ids,
+                  save_path=save_path,
+                  pl2seed_radius=pl2seed_radius,
+                  attr_tokenizer=attr_tokenizer,
+                  color_type='state',
+                  )
+
+
+if __name__ == "__main__":
+    parser = ArgumentParser()
+    parser.add_argument('--data_path', type=str, default='/u/xiuyu/work/dev4/data/waymo_processed')
+    parser.add_argument('--tfrecord_dir', type=str, default='validation_tfrecords_splitted')
+    # plot tokenized data
+    parser.add_argument('--save_folder', type=str, default='plot_gt')
+    parser.add_argument('--split', type=str, default='validation')
+    parser.add_argument('--scenario_id', type=str, default=None)
+    parser.add_argument('--plot_tokenize', action='store_true')
+    # plot generated rollouts
+    parser.add_argument('--plot_file', action='store_true')
+    parser.add_argument('--folder_path', type=str, default=None)
+    parser.add_argument('--file_path', type=str, default=None)
+    args = parser.parse_args()
+
+    if args.plot_tokenize:
+
+        scenario_id = "74ad7b76d5906d39"
+        # scenario_id = "1d60300bc06f4801"
+        data_path = os.path.join(args.data_path, args.split, f"{scenario_id}.pkl")
+        data = pickle.load(open(data_path, "rb"))
+        data['tfrecord_path'] = os.path.join(args.tfrecord_dir, f'{scenario_id}.tfrecords')
+        CONSOLE.log(f"Loaded scenario {scenario_id}")
+
+        save_path = os.path.join(args.data_path, args.save_folder, args.split)
+        os.makedirs(save_path, exist_ok=True)
+
+        plot_tokenize(data, save_path)
+
+    if args.plot_file:
+
+        plot_file(args.data_path, folder=args.folder_path, files=args.file_path)

seq_10_150_3_3_encode_occ_separate_offsets_bs8_128_no_seqindex_long_lastvalid_ablation_grid/backups/environment.yml

@@ -0,0 +1,326 @@
+name: traj
+channels:
+  - conda-forge
+  - defaults
+dependencies:
+  - _libgcc_mutex=0.1=conda_forge
+  - _openmp_mutex=4.5=2_gnu
+  - ca-certificates=2025.1.31=hbcca054_0
+  - ld_impl_linux-64=2.38=h1181459_1
+  - libffi=3.4.4=h6a678d5_1
+  - libgcc=14.2.0=h77fa898_1
+  - libgcc-ng=14.2.0=h69a702a_1
+  - libgomp=14.2.0=h77fa898_1
+  - libstdcxx-ng=11.2.0=h1234567_1
+  - ncdu=1.16=h0f457ee_0
+  - ncurses=6.4=h6a678d5_0
+  - openssl=3.4.0=h7b32b05_1
+  - pip=24.2=py39h06a4308_0
+  - python=3.9.19=h955ad1f_1
+  - readline=8.2=h5eee18b_0
+  - sqlite=3.45.3=h5eee18b_0
+  - tk=8.6.14=h39e8969_0
+  - wheel=0.43.0=py39h06a4308_0
+  - xz=5.4.6=h5eee18b_1
+  - zlib=1.2.13=h5eee18b_1
+  - pip:
+      - absl-py==1.4.0
+      - addict==2.4.0
+      - aiohappyeyeballs==2.4.0
+      - aiohttp==3.10.5
+      - aiosignal==1.3.1
+      - anyio==4.4.0
+      - appdirs==1.4.4
+      - argon2-cffi==23.1.0
+      - argon2-cffi-bindings==21.2.0
+      - array-record==0.5.1
+      - arrow==1.3.0
+      - asttokens==2.4.1
+      - astunparse==1.6.3
+      - async-lru==2.0.4
+      - async-timeout==4.0.3
+      - attrs==24.2.0
+      - av==12.3.0
+      - babel==2.16.0
+      - beautifulsoup4==4.12.3
+      - bidict==0.23.1
+      - bleach==6.1.0
+      - blinker==1.8.2
+      - cachetools==5.5.0
+      - certifi==2024.7.4
+      - cffi==1.17.0
+      - chardet==5.2.0
+      - charset-normalizer==3.3.2
+      - click==8.1.7
+      - cloudpickle==3.0.0
+      - colorlog==6.8.2
+      - comet-ml==3.45.0
+      - comm==0.2.2
+      - configargparse==1.7
+      - configobj==5.0.8
+      - contourpy==1.3.0
+      - cryptography==43.0.0
+      - cycler==0.12.1
+      - dacite==1.8.1
+      - dash==2.17.1
+      - dash-core-components==2.0.0
+      - dash-html-components==2.0.0
+      - dash-table==5.0.0
+      - dask==2023.3.1
+      - dataclass-array==1.5.1
+      - debugpy==1.8.5
+      - decorator==5.1.1
+      - defusedxml==0.7.1
+      - descartes==1.1.0
+      - dm-tree==0.1.8
+      - docker-pycreds==0.4.0
+      - docstring-parser==0.16
+      - dulwich==0.22.1
+      - easydict==1.13
+      - einops==0.8.0
+      - einsum==0.3.0
+      - embreex==2.17.7.post5
+      - etils==1.5.2
+      - eval-type-backport==0.2.0
+      - everett==3.1.0
+      - exceptiongroup==1.2.2
+      - executing==2.0.1
+      - fastjsonschema==2.20.0
+      - filelock==3.15.4
+      - fire==0.6.0
+      - flask==3.0.3
+      - flatbuffers==24.3.25
+      - fonttools==4.53.1
+      - fqdn==1.5.1
+      - frozenlist==1.4.1
+      - fsspec==2024.6.1
+      - gast==0.4.0
+      - gdown==5.2.0
+      - gitdb==4.0.11
+      - gitpython==3.1.43
+      - google-auth==2.16.2
+      - google-auth-oauthlib==1.0.0
+      - google-pasta==0.2.0
+      - grpcio==1.66.1
+      - h11==0.14.0
+      - h5py==3.11.0
+      - httpcore==1.0.5
+      - httpx==0.27.2
+      - idna==3.8
+      - imageio==2.35.1
+      - immutabledict==2.2.0
+      - importlib-metadata==8.4.0
+      - importlib-resources==6.4.4
+      - ipykernel==6.29.5
+      - ipython==8.18.1
+      - ipywidgets==8.1.5
+      - isoduration==20.11.0
+      - itsdangerous==2.2.0
+      - jax==0.4.30
+      - jaxlib==0.4.30
+      - jaxtyping==0.2.33
+      - jedi==0.19.1
+      - jinja2==3.1.4
+      - joblib==1.4.2
+      - json5==0.9.25
+      - jsonpointer==3.0.0
+      - jsonschema==4.23.0
+      - jsonschema-specifications==2023.12.1
+      - jupyter-client==8.6.2
+      - jupyter-core==5.7.2
+      - jupyter-events==0.10.0
+      - jupyter-lsp==2.2.5
+      - jupyter-server==2.14.2
+      - jupyter-server-terminals==0.5.3
+      - jupyterlab==4.2.5
+      - jupyterlab-pygments==0.3.0
+      - jupyterlab-server==2.27.3
+      - jupyterlab-widgets==3.0.13
+      - keras==2.12.0
+      - kiwisolver==1.4.5
+      - lark==1.2.2
+      - lazy-loader==0.4
+      - libclang==18.1.1
+      - lightning-utilities==0.11.6
+      - locket==1.0.0
+      - lxml==5.3.0
+      - manifold3d==2.5.1
+      - mapbox-earcut==1.0.2
+      - markdown==3.7
+      - markdown-it-py==3.0.0
+      - markupsafe==2.1.5
+      - matplotlib==3.9.2
+      - matplotlib-inline==0.1.7
+      - mdurl==0.1.2
+      - mediapy==1.2.2
+      - mistune==3.0.2
+      - ml-dtypes==0.4.0
+      - mpmath==1.3.0
+      - msgpack==1.0.8
+      - msgpack-numpy==0.4.8
+      - multidict==6.0.5
+      - namex==0.0.8
+      - nbclient==0.10.0
+      - nbconvert==7.16.4
+      - nbformat==5.10.4
+      - nerfacc==0.5.2
+      - nerfstudio==0.3.4
+      - nest-asyncio==1.6.0
+      - networkx==3.2.1
+      - ninja==1.11.1.1
+      - nodeenv==1.9.1
+      - notebook-shim==0.2.4
+      - numpy==1.23.0
+      - nuscenes-devkit==1.1.11
+      - nvidia-cublas-cu12==12.1.3.1
+      - nvidia-cuda-cupti-cu12==12.1.105
+      - nvidia-cuda-nvrtc-cu12==12.1.105
+      - nvidia-cuda-runtime-cu12==12.1.105
+      - nvidia-cudnn-cu12==9.1.0.70
+      - nvidia-cufft-cu12==11.0.2.54
+      - nvidia-curand-cu12==10.3.2.106
+      - nvidia-cusolver-cu12==11.4.5.107
+      - nvidia-cusparse-cu12==12.1.0.106
+      - nvidia-nccl-cu12==2.20.5
+      - nvidia-nvjitlink-cu12==12.6.20
+      - nvidia-nvtx-cu12==12.1.105
+      - oauthlib==3.2.2
+      - open3d==0.18.0
+      - opencv-python==4.6.0.66
+      - openexr==1.3.9
+      - opt-einsum==3.3.0
+      - optree==0.12.1
+      - overrides==7.7.0
+      - packaging==24.1
+      - pandas==1.5.3
+      - pandocfilters==1.5.1
+      - parso==0.8.4
+      - partd==1.4.2
+      - pexpect==4.9.0
+      - pillow==9.2.0
+      - platformdirs==4.2.2
+      - plotly==5.13.1
+      - prometheus-client==0.20.0
+      - promise==2.3
+      - prompt-toolkit==3.0.47
+      - protobuf==3.20.3
+      - psutil==6.0.0
+      - ptyprocess==0.7.0
+      - pure-eval==0.2.3
+      - pyarrow==10.0.0
+      - pyasn1==0.6.0
+      - pyasn1-modules==0.4.0
+      - pycocotools==2.0.8
+      - pycollada==0.8
+      - pycparser==2.22
+      - pygments==2.18.0
+      - pyliblzfse==0.4.1
+      - pymeshlab==2023.12.post1
+      - pyngrok==7.2.0
+      - pyparsing==3.1.4
+      - pyquaternion==0.9.9
+      - pysocks==1.7.1
+      - python-box==6.1.0
+      - python-dateutil==2.9.0.post0
+      - python-engineio==4.9.1
+      - python-json-logger==2.0.7
+      - python-socketio==5.11.3
+      - pytorch-lightning==2.4.0
+      - pytz==2024.1
+      - pywavelets==1.6.0
+      - pyyaml==6.0.2
+      - pyzmq==26.2.0
+      - rawpy==0.22.0
+      - referencing==0.35.1
+      - requests==2.32.3
+      - requests-oauthlib==2.0.0
+      - requests-toolbelt==1.0.0
+      - retrying==1.3.4
+      - rfc3339-validator==0.1.4
+      - rfc3986-validator==0.1.1
+      - rich==13.8.0
+      - rpds-py==0.20.0
+      - rsa==4.9
+      - rtree==1.3.0
+      - scikit-image==0.20.0
+      - scikit-learn==1.2.2
+      - scipy==1.9.1
+      - seaborn==0.13.2
+      - semantic-version==2.10.0
+      - send2trash==1.8.3
+      - sentry-sdk==2.13.0
+      - setproctitle==1.3.3
+      - setuptools==67.6.0
+      - shapely==1.8.5.post1
+      - shtab==1.7.1
+      - simple-websocket==1.0.0
+      - simplejson==3.19.3
+      - six==1.16.0
+      - smmap==5.0.1
+      - sniffio==1.3.1
+      - soupsieve==2.6
+      - splines==0.3.0
+      - stack-data==0.6.3
+      - svg-path==6.3
+      - sympy==1.13.2
+      - tenacity==9.0.0
+      - tensorboard==2.12.3
+      - tensorboard-data-server==0.7.2
+      - tensorflow==2.12.0
+      - tensorflow-addons==0.23.0
+      - tensorflow-datasets==4.9.3
+      - tensorflow-estimator==2.12.0
+      - tensorflow-graphics==2021.12.3
+      - tensorflow-io-gcs-filesystem==0.37.1
+      - tensorflow-metadata==1.15.0
+      - tensorflow-probability==0.19.0
+      - termcolor==2.4.0
+      - terminado==0.18.1
+      - threadpoolctl==3.5.0
+      - tifffile==2024.8.28
+      - timm==0.6.7
+      - tinycss2==1.3.0
+      - toml==0.10.2
+      - tomli==2.0.1
+      - toolz==0.12.1
+      - torch==2.4.0
+      - torch-cluster==1.6.3+pt24cu121
+      - torch-fidelity==0.3.0
+      - torch-geometric==2.5.3
+      - torch-scatter==2.1.2+pt24cu121
+      - torch-sparse==0.6.18+pt24cu121
+      - torchmetrics==1.4.1
+      - torchvision==0.19.0
+      - tornado==6.4.1
+      - tqdm==4.66.5
+      - traitlets==5.14.3
+      - trimesh==4.4.7
+      - triton==3.0.0
+      - typeguard==2.13.3
+      - types-python-dateutil==2.9.0.20240821
+      - typing-extensions==4.12.2
+      - tyro==0.8.10
+      - tzdata==2024.1
+      - uri-template==1.3.0
+      - urllib3==2.2.2
+      - vhacdx==0.0.8.post1
+      - viser==0.1.3
+      - visu3d==1.5.1
+      - wandb==0.17.8
+      - waymo-open-dataset-tf-2-12-0==1.6.4
+      - wcwidth==0.2.13
+      - webcolors==24.8.0
+      - webencodings==0.5.1
+      - websocket-client==1.8.0
+      - websockets==13.0.1
+      - werkzeug==3.0.4
+      - widgetsnbextension==4.0.13
+      - wrapt==1.14.1
+      - wsproto==1.2.0
+      - wurlitzer==3.1.1
+      - xatlas==0.0.9
+      - xxhash==3.5.0
+      - yarl==1.9.11
+      - yourdfpy==0.0.56
+      - zipp==3.20.1
+prefix: /u/xiuyu/anaconda3/envs/traffic

seq_10_150_3_3_encode_occ_separate_offsets_bs8_128_no_seqindex_long_lastvalid_ablation_grid/backups/run.py

@@ -0,0 +1,181 @@
+import pytorch_lightning as pl
+import os
+import shutil
+import fnmatch
+import torch
+from argparse import ArgumentParser
+from pytorch_lightning.callbacks import LearningRateMonitor
+from pytorch_lightning.callbacks import ModelCheckpoint
+from pytorch_lightning.strategies import DDPStrategy
+from pytorch_lightning.loggers import WandbLogger
+
+from dev.utils.func import RankedLogger, load_config_act, CONSOLE
+from dev.datasets.scalable_dataset import MultiDataModule
+from dev.model.smart import SMART
+
+
+def backup(source_dir, backup_dir):
+    """
+    Back up the source directory (code and configs) to a backup directory.
+    """
+
+    if os.path.exists(backup_dir):
+        return
+    os.makedirs(backup_dir, exist_ok=False)
+
+    # Helper function to check if a path matches exclude patterns
+    def should_exclude(path):
+        for pattern in exclude_patterns:
+            if fnmatch.fnmatch(os.path.basename(path), pattern):
+                return True
+        return False
+
+    # Iterate through the files and directories in source_dir
+    for root, dirs, files in os.walk(source_dir):
+        # Skip excluded directories
+        dirs[:] = [d for d in dirs if not should_exclude(d)]
+
+        # Determine the relative path and destination path
+        rel_path = os.path.relpath(root, source_dir)
+        dest_dir = os.path.join(backup_dir, rel_path)
+        os.makedirs(dest_dir, exist_ok=True)
+        
+        # Copy all relevant files
+        for file in files:
+            if any(fnmatch.fnmatch(file, pattern) for pattern in include_patterns):
+                shutil.copy2(os.path.join(root, file), os.path.join(dest_dir, file))
+    
+    logger.info(f"Backup completed. Files saved to: {backup_dir}")
+
+
+if __name__ == '__main__':
+    pl.seed_everything(2024, workers=True)
+    torch.set_printoptions(precision=3)
+
+    parser = ArgumentParser()
+    parser.add_argument('--config', type=str, default='configs/ours_long_term.yaml')
+    parser.add_argument('--pretrain_ckpt', type=str, default=None,
+                        help='Path to any pretrained model, will only load its parameters.'
+    )
+    parser.add_argument('--ckpt_path', type=str, default=None,
+                        help='Path to any trained model, will load all the states.'
+    )
+    parser.add_argument('--save_ckpt_path', type=str, default='output/debug',
+                        help='Path to save the checkpoints in training mode'
+    )
+    parser.add_argument('--save_path', type=str, default=None,
+                        help='Path to save the inference results in validation and test mode.'
+    )
+    parser.add_argument('--wandb', action='store_true',
+                        help='Whether to use wandb logger in training.'
+    )
+    parser.add_argument('--devices', type=int, default=1)
+    parser.add_argument('--train', action='store_true')
+    parser.add_argument('--validate', action='store_true')
+    parser.add_argument('--test', action='store_true')
+    parser.add_argument('--plot_rollouts', action='store_true')
+    args = parser.parse_args()
+
+    if not (args.train or args.validate or args.test or args.plot_rollouts):
+        raise RuntimeError(f"Got invalid action, should be one of ['train', 'validate', 'test', 'plot_rollouts']")
+
+    # ! setup logger
+    logger = RankedLogger(__name__, rank_zero_only=True)
+
+    # ! backup codes
+    exclude_patterns = ['*output*', '*logs', 'wandb', 'data', '*debug*', '*backup*', 'interact_*', '*edge_map*', '__pycache__']
+    include_patterns = ['*.py', '*.json', '*.yaml', '*.yml', '*.sh']
+    backup(os.getcwd(), os.path.join(args.save_ckpt_path, 'backups'))
+
+    config = load_config_act(args.config)
+
+    wandb_logger = None
+    if args.wandb and not int(os.getenv('DEBUG', 0)):
+        # squeue -O username,state,nodelist,gres,minmemory,numcpus,name
+        wandb_logger = WandbLogger(project='simagent')
+
+    trainer_config = config.Trainer
+    max_epochs = trainer_config.max_epochs
+
+    # ! setup datamodule and model
+    datamodule = MultiDataModule(**vars(config.Dataset), logger=logger)
+    model = SMART(config.Model, save_path=args.save_ckpt_path, logger=logger, max_epochs=max_epochs)
+    if args.pretrain_ckpt:
+        model.load_state_from_file(filename=args.pretrain_ckpt)
+    strategy = DDPStrategy(find_unused_parameters=True, gradient_as_bucket_view=True)
+    logger.info(f'Build model: {model.__class__.__name__} datamodule: {datamodule.__class__.__name__}')
+
+    # ! checkpoint configuration
+    every_n_epochs = 1
+    if int(os.getenv('OVERFIT', 0)):
+        max_epochs = trainer_config.overfit_epochs
+        every_n_epochs = 100
+
+    if int(os.getenv('CHECK_INPUTS', 0)):
+        max_epochs = 1
+
+    check_val_every_n_epoch = 1  # save checkpoints for each epoch
+    model_checkpoint = ModelCheckpoint(dirpath=args.save_ckpt_path,
+                                       filename='{epoch:02d}',
+                                       save_top_k=5,
+                                       monitor='epoch',
+                                       mode='max',
+                                       save_last=True,
+                                       every_n_train_steps=1000,
+                                       save_on_train_epoch_end=True)
+
+    # ! setup trainer
+    lr_monitor = LearningRateMonitor(logging_interval='epoch')
+    trainer = pl.Trainer(accelerator=trainer_config.accelerator, devices=args.devices if args.devices is not None else trainer_config.devices,
+                         strategy=strategy, logger=wandb_logger,
+                         accumulate_grad_batches=trainer_config.accumulate_grad_batches,
+                         num_nodes=trainer_config.num_nodes,
+                         callbacks=[model_checkpoint, lr_monitor],
+                         max_epochs=max_epochs,
+                         num_sanity_val_steps=0,
+                         check_val_every_n_epoch=check_val_every_n_epoch,
+                         log_every_n_steps=1,
+                         gradient_clip_val=0.5)
+    logger.info(f'Build trainer: {trainer.__class__.__name__}')
+
+    # ! run
+    if args.train:
+
+        logger.info(f'Start training ...')
+        trainer.fit(model, datamodule, ckpt_path=args.ckpt_path)
+
+    # NOTE: here both validation and test process use validation split data
+    # for validation, we enable the online metric calculation with results dumping
+    # for test, we disable it and only dump the inference results.
+    else:
+
+        if args.save_path is not None:
+            save_path = args.save_path
+        else:
+            assert args.ckpt_path is not None and os.path.exists(args.ckpt_path), \
+                    f'Path {args.ckpt_path} not exists!'
+            save_path = os.path.join(os.path.dirname(args.ckpt_path), 'validation')
+        os.makedirs(save_path, exist_ok=True)
+        CONSOLE.log(f'Results will be saved to [yellow]{save_path}[/]')
+
+        model.save_path = save_path
+
+        if not args.ckpt_path:
+            CONSOLE.log(f'[yellow] Warning: no checkpoint will be loaded in validation! [/]')
+
+        if args.validate:
+
+            CONSOLE.log('[on blue] Start validating ... [/]')
+            model.set(mode='validation')
+
+        elif args.test:
+
+            CONSOLE.log('[on blue] Sart testing ... [/]')
+            model.set(mode='test')
+
+        elif args.plot_rollouts:
+
+            CONSOLE.log('[on blue] Sart generating ... [/]')
+            model.set(mode='plot_rollouts')
+
+        trainer.validate(model, datamodule, ckpt_path=args.ckpt_path)

seq_10_150_3_3_encode_occ_separate_offsets_bs8_128_no_seqindex_long_lastvalid_ablation_grid/backups/scripts/aggregate_log_metric_features.sh

@@ -0,0 +1,16 @@
+#!/bin/bash
+
+export TF_CPP_MIN_LOG_LEVEL='2'
+export PYTHONPATH='.'
+
+# dump all features
+echo 'Start dump all log features ...'
+python dev/metrics/compute_metrics.py --dump_log --no_batch
+
+sleep 20
+
+# aggregate features
+echo 'Start aggregate log features ...'
+python dev/metrics/compute_metrics.py --aggregate_log
+
+echo 'Done!

seq_10_150_3_3_encode_occ_separate_offsets_bs8_128_no_seqindex_long_lastvalid_ablation_grid/backups/scripts/c128.sh

@@ -0,0 +1,13 @@
+#!/bin/bash
+
+#SBATCH --job-name c128                   # Job name
+### Logging
+#SBATCH --output=%j.out                 # Stdout (%j expands to jobId)
+#SBATCH --error=%j.err                  # Stderr (%j expands to jobId)
+### Node info
+#SBATCH --nodes=1                       # Single node or multi node
+#SBATCH --time 100:00:00                 # Max time (hh:mm:ss)
+#SBATCH --gres=gpu:0                    # GPUs per node
+#SBATCH --mem=128G                      # Recommend 32G per GPU
+#SBATCH --ntasks-per-node=1             # Tasks per node
+#SBATCH --cpus-per-task=64              # Recommend 8 per GPU

seq_10_150_3_3_encode_occ_separate_offsets_bs8_128_no_seqindex_long_lastvalid_ablation_grid/backups/scripts/c64.sh

@@ -0,0 +1,13 @@
+#!/bin/bash
+
+#SBATCH --job-name c128                   # Job name
+### Logging
+#SBATCH --output=%j.out                 # Stdout (%j expands to jobId)
+#SBATCH --error=%j.err                  # Stderr (%j expands to jobId)
+### Node info
+#SBATCH --nodes=1                       # Single node or multi node
+#SBATCH --time 100:00:00                 # Max time (hh:mm:ss)
+#SBATCH --gres=gpu:0                    # GPUs per node
+#SBATCH --mem=128G                      # Recommend 32G per GPU
+#SBATCH --ntasks-per-node=1             # Tasks per node
+#SBATCH --cpus-per-task=64              # Recommend 8 per GPU

seq_10_150_3_3_encode_occ_separate_offsets_bs8_128_no_seqindex_long_lastvalid_ablation_grid/backups/scripts/compute_metrics.sh

@@ -0,0 +1,13 @@
+#!/bin/bash
+
+export TORCH_LOGS='0'
+export TF_CPP_MIN_LOG_LEVEL='2'
+export PYTHONPATH='.'
+
+NUM_WORKERS=$1
+SIM_DIR=$2
+
+echo 'Start running ...'
+python dev/metrics/compute_metrics.py --compute_metric --num_workers "$NUM_WORKERS" --sim_dir "$SIM_DIR" ${@:3}
+
+echo 'Done!

seq_10_150_3_3_encode_occ_separate_offsets_bs8_128_no_seqindex_long_lastvalid_ablation_grid/backups/scripts/data_preprocess.sh

@@ -0,0 +1,12 @@
+#!/bin/bash
+
+# env
+source ~/anaconda3/etc/profile.d/conda.sh
+conda config --append envs_dirs ~/.conda/envs
+conda activate traj
+
+echo "Starting running..."
+
+# multi-GPU training
+cd ~/work/dev6/thirdparty/dev4
+PYTHONPATH='..':$PYTHONPATH python3 data_preprocess.py --split validation

seq_10_150_3_3_encode_occ_separate_offsets_bs8_128_no_seqindex_long_lastvalid_ablation_grid/backups/scripts/data_preprocess_loop.sh

@@ -0,0 +1,23 @@
+#!/bin/bash
+
+RED='\033[0;31m'
+NC='\033[0m'
+
+cd /u/xiuyu/work/dev4/
+
+trap "echo -e \"${RED}Stopping script...${NC}\"; kill -- -$$" SIGINT
+
+while true; do
+    echo -e "${RED}Start running ...${NC}"
+    PYTHONPATH='.':$PYTHONPATH setsid python data_preprocess.py --split training &
+    PID=$!
+
+    sleep 1200
+
+    echo -e "${RED}Sending SIGINT to process group $PID...${NC}"
+    PGID=$(ps -o pgid= -p $PID | tail -n 1 | tr -d ' ')
+    kill -- -$PGID
+    wait $PID
+
+    sleep 10
+done

seq_10_150_3_3_encode_occ_separate_offsets_bs8_128_no_seqindex_long_lastvalid_ablation_grid/backups/scripts/debug.py

@@ -0,0 +1,17 @@
+import torch
+from torch import nn
+import torch.optim as optim
+
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+
+embedding = nn.Embedding(180, 128).to(device)
+gt = torch.randint(0, 2, (180, 2048)).to(device)
+head = nn.Linear(128, 2048).to(device)
+optimizer = optim.Adam([embedding.weight, head.weight])
+
+while True:
+    pred = head(embedding.weight).sigmoid()
+    loss = nn.MSELoss()(pred, gt.float())
+    optimizer.zero_grad()
+    loss.backward()
+    optimizer.step()

seq_10_150_3_3_encode_occ_separate_offsets_bs8_128_no_seqindex_long_lastvalid_ablation_grid/backups/scripts/debug_map.py

@@ -0,0 +1,204 @@
+import os
+import pickle
+import torch
+import matplotlib.pyplot as plt
+import numpy as np
+from tqdm import tqdm
+from argparse import ArgumentParser
+from dev.datasets.preprocess import TokenProcessor
+from dev.transforms.target_builder import WaymoTargetBuilder
+
+
+colors = [
+    ('#1f77b4', '#1a5a8a'),  # blue
+    ('#2ca02c', '#217721'),  # green
+    ('#ff7f0e', '#cc660b'),  # orange
+    ('#9467bd', '#6f4a91'),  # purple
+    ('#d62728', '#a31d1d'),  # red
+    ('#000000', '#000000'),  # black
+]
+
+
+def draw_map(tokenize_data, token_processor: TokenProcessor, index, posfix):
+    print("Drawing raw data ...")
+    shift = 5
+    token_size = 2048
+
+    traj_token = token_processor.trajectory_token["veh"]
+    traj_token_all = token_processor.trajectory_token_all["veh"]
+
+    plt.subplots_adjust(left=0.3, right=0.7, top=0.7, bottom=0.3)
+    fig, ax = plt.subplots()
+    ax.set_axis_off()
+
+    scenario_id = data['scenario_id']
+    ax.scatter(tokenize_data["map_point"]["position"][:, 0],
+               tokenize_data["map_point"]["position"][:, 1], s=0.2, c='black', edgecolors='none')
+
+    index = np.array(index).astype(np.int32)
+    agent_data = tokenize_data["agent"]
+    token_index = agent_data["token_idx"][index]
+    token_valid_mask = agent_data["agent_valid_mask"][index]
+
+    num_agent, num_token = token_index.shape
+    tokens = traj_token[token_index.view(-1)].reshape(num_agent, num_token, 4, 2)
+    tokens_all = traj_token_all[token_index.view(-1)].reshape(num_agent, num_token, 6, 4, 2)
+
+    position = agent_data['position'][index, :, :2] # (num_agent, 91, 2)
+    heading = agent_data['heading'][index] # (num_agent, 91)
+    valid_mask = (position[..., 0] != 0) & (position[..., 1] != 0) # (num_agent, 91)
+    # TODO: fix this
+    if args.smart:
+        for shifted_tid in range(token_valid_mask.shape[1]):
+            valid_mask[:, shifted_tid * shift : (shifted_tid + 1) * shift] = token_valid_mask[:, shifted_tid : shifted_tid + 1].repeat(1, shift)
+    else:
+        for shifted_tid in range(token_index.shape[1]):
+            valid_mask[:, shifted_tid * shift : (shifted_tid + 1) * shift] = token_index[:, shifted_tid : shifted_tid + 1] != token_size + 2
+    last_valid_step = valid_mask.shape[1] - 1 - torch.argmax(valid_mask.flip(dims=[1]).long(), dim=1)
+    last_valid_step = {int(index[i]): int(last_valid_step[i]) for i in range(len(index))}
+
+    _, token_num, token_contour_dim, feat_dim = tokens.shape
+    tokens_src = tokens.reshape(num_agent, token_num * token_contour_dim, feat_dim)
+    tokens_all_src = tokens_all.reshape(num_agent, token_num * 6 * token_contour_dim, feat_dim)
+    prev_heading = heading[:, 0]
+    prev_pos = position[:, 0]
+
+    fig_paths = []
+    agent_colors = np.zeros((num_agent, position.shape[1]))
+    shape = np.zeros((num_agent, position.shape[1], 2)) + 3.
+    for tid in tqdm(range(shift, position.shape[1], shift), leave=False, desc="Token ..."):
+        cos, sin = prev_heading.cos(), prev_heading.sin()
+        rot_mat = prev_heading.new_zeros(num_agent, 2, 2)
+        rot_mat[:, 0, 0] = cos
+        rot_mat[:, 0, 1] = sin
+        rot_mat[:, 1, 0] = -sin
+        rot_mat[:, 1, 1] = cos
+        tokens_world = torch.bmm(torch.from_numpy(tokens_src).float(), rot_mat).reshape(num_agent,
+                                                                                        token_num,
+                                                                                        token_contour_dim,
+                                                                                        feat_dim)
+        tokens_all_world = torch.bmm(torch.from_numpy(tokens_all_src).float(), rot_mat).reshape(num_agent,
+                                                                                                token_num,
+                                                                                                6,
+                                                                                                token_contour_dim,
+                                                                                                feat_dim)
+        tokens_world += prev_pos[:, None, None, :2]
+        tokens_all_world += prev_pos[:, None, None, None, :2]
+        tokens_select = tokens_world[:, tid // shift - 1] # (num_agent, token_contour_dim, feat_dim)
+        tokens_all_select = tokens_all_world[:, tid // shift - 1] # (num_agent, 6, token_contour_dim, feat_dim)
+
+        diff_xy = tokens_select[:, 0, :] - tokens_select[:, 3, :]
+        prev_heading = heading[:, tid].clone()
+        # prev_heading[valid_mask[:, tid - shift]] = torch.arctan2(diff_xy[:, 1], diff_xy[:, 0])[
+        #     valid_mask[:, tid - shift]]
+        prev_pos = position[:, tid].clone()
+        # prev_pos[valid_mask[:, tid - shift]] = tokens_select.mean(dim=1)[valid_mask[:, tid - shift]]
+
+        # NOTE tokens_pos equals to tokens_all_pos[:, -1]
+        tokens_pos = tokens_select.mean(dim=1) # (num_agent, 2)
+        tokens_all_pos = tokens_all_select.mean(dim=2) # (num_agent, 6, 2)
+
+        # colors
+        cur_token_index = token_index[:, tid // shift - 1]
+        is_bos = cur_token_index == token_size
+        is_eos = cur_token_index == token_size + 1
+        is_invalid = cur_token_index == token_size + 2
+        is_valid = ~is_bos & ~is_eos & ~is_invalid
+        agent_colors[is_valid, tid - shift : tid] = 1
+        agent_colors[is_bos, tid - shift : tid] = 2
+        agent_colors[is_eos, tid - shift : tid] = 3
+        agent_colors[is_invalid, tid - shift : tid] = 4
+
+        for i in tqdm(range(shift), leave=False, desc="Timestep ..."):
+            global_tid = tid - shift + i
+            cur_valid_mask = valid_mask[:, tid - shift] # only when the last tokenized timestep is valid the current shifts trajectory is valid
+            xs = tokens_all_pos[cur_valid_mask, i, 0]
+            ys = tokens_all_pos[cur_valid_mask, i, 1]
+            widths = shape[cur_valid_mask, global_tid, 1]
+            lengths = shape[cur_valid_mask, global_tid, 0]
+            angles = heading[cur_valid_mask, global_tid]
+            cur_agent_colors = agent_colors[cur_valid_mask, global_tid]
+            current_index = index[cur_valid_mask]
+
+            drawn_agents = []
+            drawn_texts = []
+            for x, y, width, length, angle, color_type, id in zip(
+                xs, ys, widths, lengths, angles, cur_agent_colors, current_index):
+                if x < 3000: continue
+                agent = plt.Rectangle((x, y), width, length, # angle=((angle + np.pi / 2) / np.pi * 360) % 360,
+                                        linewidth=0.2,
+                                        facecolor=colors[int(color_type) - 1][0],
+                                        edgecolor=colors[int(color_type) - 1][1])
+                ax.add_patch(agent)
+                text = plt.text(x-4, y-4, f"{str(id)}:{str(global_tid)}", fontdict={'family': 'serif', 'size': 3, 'color': 'red'})
+
+                if global_tid != last_valid_step[id]:
+                    drawn_agents.append(agent)
+                    drawn_texts.append(text)
+
+                # draw timestep to be tokenized
+                if global_tid % shift == 0:
+                    tokenize_agent = plt.Rectangle((x, y), width, length, # angle=((angle + np.pi / 2) / np.pi * 360) % 360,
+                                                    linewidth=0.2, fill=False,
+                                                    edgecolor=colors[int(color_type) - 1][1])
+                    ax.add_patch(tokenize_agent)
+
+            plt.gca().set_aspect('equal', adjustable='box')
+
+            fig_path = f"debug/tokenize/steps/{scenario_id}_{global_tid}.png"
+            plt.savefig(fig_path, dpi=600, bbox_inches="tight")
+            fig_paths.append(fig_path)
+
+            for drawn_agent, drawn_text in zip(drawn_agents, drawn_texts):
+                drawn_agent.remove()
+                drawn_text.remove()
+
+    plt.close()
+
+    # generate gif
+    import imageio.v2 as imageio
+    images = []
+    for fig_path in tqdm(fig_paths, leave=False, desc="Generate gif ..."):
+        images.append(imageio.imread(fig_path))
+    imageio.mimsave(f"debug/tokenize/{scenario_id}_tokenize_{posfix}.gif", images, duration=0.1)
+
+
+def main(data):
+
+    token_size = 2048
+
+    os.makedirs("debug/tokenize/steps/", exist_ok=True)
+    scenario_id = data["scenario_id"]
+
+    selected_agents_index = [1, 21, 35, 36, 46]
+
+    # raw data
+    if not os.path.exists(f"debug/tokenize/{scenario_id}_raw.gif"):
+        draw_raw(data, selected_agents_index)
+
+    # tokenization
+    token_processor = TokenProcessor(token_size, disable_invalid=args.smart)
+    print(f"Loaded token processor with token_size: {token_size}")
+    data = token_processor.preprocess(data)
+
+    # tokenzied data
+    posfix = "smart" if args.smart else "ours"
+    # if not os.path.exists(f"debug/tokenize/{scenario_id}_tokenize_{posfix}.gif"):
+    draw_tokenize(data, token_processor, selected_agents_index, posfix)
+
+    target_builder = WaymoTargetBuilder(num_historical_steps=11, num_future_steps=80)
+    data = target_builder(data)
+
+
+if __name__ == "__main__":
+    parser = ArgumentParser(description="Testing script parameters")
+    parser.add_argument("--smart", action="store_true")
+    parser.add_argument("--data_path", type=str, default="/u/xiuyu/work/dev4/data/waymo_processed/training")
+    args = parser.parse_args()
+
+    scenario_id = "74ad7b76d5906d39"
+    data_path = os.path.join(args.data_path, f"{scenario_id}.pkl")
+    data = pickle.load(open(data_path, "rb"))
+    print(f"Loaded scenario {scenario_id}")
+
+    main(data)

seq_10_150_3_3_encode_occ_separate_offsets_bs8_128_no_seqindex_long_lastvalid_ablation_grid/backups/scripts/g2.sh

@@ -0,0 +1,37 @@
+#!/bin/bash
+
+#SBATCH --job-name g2                   # Job name
+### Logging
+#SBATCH --output=%j.out                 # Stdout (%j expands to jobId)
+#SBATCH --error=%j.err                  # Stderr (%j expands to jobId)
+### Node info
+#SBATCH --nodes=1                       # Single node or multi node
+#SBATCH --nodelist=sota-2
+#SBATCH --time 24:00:00                 # Max time (hh:mm:ss)
+#SBATCH --gres=gpu:2                    # GPUs per node
+#SBATCH --mem=96G                      # Recommend 32G per GPU
+#SBATCH --ntasks-per-node=1             # Tasks per node
+#SBATCH --cpus-per-task=16              # Recommend 8 per GPU
+
+export NCCL_DEBUG=INFO
+export REQUESTS_CA_BUNDLE="/etc/ssl/certs/ca-certificates.crt"
+export HTTPS_PROXY="https://192.168.0.10:443/"
+export https_proxy="https://192.168.0.10:443/"
+
+export TEST_VAL_TRAIN=False
+export TEST_VAL_PRED=True
+export WANDB=True
+
+sleep 86400
+
+cd /u/xiuyu/work/dev4
+PYTHONPATH=".":$PYTHONPATH python3 train.py \
+                            --devices 2 \
+                            --config configs/train/train_scalable_with_state.yaml \
+                            --save_ckpt_path output/seed_1k_pure_seed_150_3_emb_head_3_debug \
+                            --pretrain_ckpt output/ours_map_pretrain/epoch=31.ckpt
+
+PYTHONPATH=".":$PYTHONPATH python val.py \
+                            --config configs/validation/val_scalable_with_state.yaml \
+                            --save_path output/seed_debug \
+                            --pretrain_ckpt output/seed_1k_pure_seed_150_3_emb_head_3/last.ckpt

seq_10_150_3_3_encode_occ_separate_offsets_bs8_128_no_seqindex_long_lastvalid_ablation_grid/backups/scripts/g4.sh

@@ -0,0 +1,35 @@
+#!/bin/bash
+
+#SBATCH --job-name g4                   # Job name
+### Logging
+#SBATCH --output=%j.out                 # Stdout (%j expands to jobId)
+#SBATCH --error=%j.err                  # Stderr (%j expands to jobId)
+### Node info
+#SBATCH --nodes=1                       # Single node or multi node
+#SBATCH --nodelist=sota-1
+#SBATCH --time 72:00:00                 # Max time (hh:mm:ss)
+#SBATCH --gres=gpu:4                    # GPUs per node
+#SBATCH --mem=128G                      # Recommend 32G per GPU
+#SBATCH --ntasks-per-node=1             # Tasks per node
+#SBATCH --cpus-per-task=32              # Recommend 8 per GPU
+
+export REQUESTS_CA_BUNDLE="/etc/ssl/certs/ca-certificates.crt"
+export HTTPS_PROXY="https://192.168.0.10:443/"
+export https_proxy="https://192.168.0.10:443/"
+
+export TEST_VAL_TRAIN=0
+export TEST_VAL_PRED=1
+export WANDB=1
+
+sleep 604800
+
+cd /u/xiuyu/work/dev4
+PYTHONPATH=".":$PYTHONPATH python3 train.py \
+                            --devices 4 \
+                            --config configs/train/train_scalable_with_state.yaml \
+                            --save_ckpt_path output/seq_1k_10_150_3_3_encode_occ_separate_offsets \
+                            --pretrain_ckpt output/pretrain_scalable_map/epoch=31.ckpt
+
+PYTHONPATH=".":$PYTHONPATH python val.py \
+                            --config configs/ours_long_term.yaml \
+                            --ckpt_path output/seq_5k_10_150_3_3_encode_occ_separate_offsets_bs8_128_no_seqindex_long/epoch=31.ckpt

seq_10_150_3_3_encode_occ_separate_offsets_bs8_128_no_seqindex_long_lastvalid_ablation_grid/backups/scripts/g8.sh

@@ -0,0 +1,44 @@
+#!/bin/bash
+
+#SBATCH --job-name g8                   # Job name
+### Logging
+#SBATCH --output=%j.out                 # Stdout (%j expands to jobId)
+#SBATCH --error=%j.err                  # Stderr (%j expands to jobId)
+### Node info
+#SBATCH --nodes=1                       # Single node or multi node
+#SBATCH --nodelist=sota-6
+#SBATCH --time 120:00:00                 # Max time (hh:mm:ss)
+#SBATCH --gres=gpu:8                    # GPUs per node
+#SBATCH --mem=256G                      # Recommend 32G per GPU
+#SBATCH --ntasks-per-node=1             # Tasks per node
+#SBATCH --cpus-per-task=32              # Recommend 8 per GPU
+
+export REQUESTS_CA_BUNDLE="/etc/ssl/certs/ca-certificates.crt"
+export HTTPS_PROXY="https://192.168.0.10:443/"
+export https_proxy="https://192.168.0.10:443/"
+
+export TEST_VAL_TRAIN=0
+export TEST_VAL_PRED=1
+export WANDB=1
+
+sleep 864000
+
+cd /u/xiuyu/work/dev4
+PYTHONPATH=".":$PYTHONPATH python3 train.py \
+                            --devices 8 \
+                            --config configs/train/train_scalable_long_term.yaml \
+                            --save_ckpt_path output/seq_5k_10_150_3_3_encode_occ_separate_offsets_bs8_128_no_seqindex_long \
+                            --pretrain_ckpt output/pretrain_scalable_map/epoch=31.ckpt
+
+PYTHONPATH=".":$PYTHONPATH python3 train.py \
+                            --devices 8 \
+                            --config configs/ours_long_term.yaml \
+                            --save_ckpt_path output2/seq_10_150_3_3_encode_occ_separate_offsets_bs8_128_no_seqindex_long_lastvalid
+
+PYTHONPATH=".":$PYTHONPATH python3 train.py \
+                            --config configs/ours_long_term.yaml \
+                            --save_ckpt_path output2/debug
+
+PYTHONPATH=".":$PYTHONPATH python val.py \
+                            --config configs/ours_long_term.yaml \
+                            --ckpt_path output2/bug_seq_10_150_3_3_encode_occ_separate_offsets_bs8_128_no_seqindex_long/last.ckpt

seq_10_150_3_3_encode_occ_separate_offsets_bs8_128_no_seqindex_long_lastvalid_ablation_grid/backups/scripts/hf_model.py

@@ -0,0 +1,111 @@
+import argparse
+import os
+from huggingface_hub import upload_folder, upload_file, hf_hub_download
+from rich.console import Console
+from rich.panel import Panel
+from rich import box, style
+from rich.table import Table
+
+CONSOLE = Console(width=120)
+
+
+def upload():
+
+    if args.folder_path:
+
+        try:
+            if token is not None:
+                upload_folder(repo_id=args.repo_id, folder_path=args.folder_path, ignore_patterns=ignore_patterns, token=token)
+            else:
+                upload_folder(repo_id=args.repo_id, folder_path=args.folder_path, ignore_patterns=ignore_patterns)
+            table = Table(title=None, show_header=False, box=box.MINIMAL, title_style=style.Style(bold=True))
+            table.add_row(f"Model id {args.repo_id}", str(args.folder_path))
+            CONSOLE.print(Panel(table, title="[bold][green]:tada: Upload completed DO NOT forget specify the model id in methods! :tada:[/bold]", expand=False))
+
+        except Exception as e:
+            CONSOLE.print(f"[bold][yellow]:tada: Upload failed due to {e}.")
+            raise e
+
+    if args.file_path:
+
+        try:
+            if token is not None:
+                upload_file(
+                    path_or_fileobj=args.file_path,
+                    path_in_repo=os.path.basename(args.file_path),
+                    repo_id=args.repo_id,
+                    repo_type='model',
+                    token=token
+                )
+            else:
+                upload_file(
+                    path_or_fileobj=args.file_path,
+                    path_in_repo=os.path.basename(args.file_path),
+                    repo_id=args.repo_id,
+                    repo_type='model',
+                )
+            table = Table(title=None, show_header=False, box=box.MINIMAL, title_style=style.Style(bold=True))
+            table.add_row(f"Model id {args.repo_id}", str(args.file_path))
+            CONSOLE.print(Panel(table, title="[bold][green]:tada: Upload completed! :tada:[/bold]", expand=False))
+
+        except Exception as e:
+            CONSOLE.print(f"[bold][yellow]:tada: Upload failed due to {e}.")
+            raise e
+
+
+def download():
+
+    try:
+        if token is not None:
+            ckpt_path = hf_hub_download(
+                repo_id=args.repo_id,
+                filename=args.file_path,
+                token=token
+            )
+        else:
+            ckpt_path = hf_hub_download(
+                repo_id=args.repo_id,
+                filename=args.file_path,
+            )
+        table = Table(title=None, show_header=False, box=box.MINIMAL, title_style=style.Style(bold=True))
+        table.add_row(f"Model id {args.repo_id}", str(args.file_path))
+        CONSOLE.print(Panel(table, title=f"[bold][green]:tada: Download completed to {ckpt_path}! :tada:[/bold]", expand=False))
+
+        if args.save_path is not None:
+            os.makedirs(args.save_path, exist_ok=True)
+            import shutil
+            shutil.copy(ckpt_path, os.path.join(args.save_path, args.file_path))
+
+    except Exception as e:
+        CONSOLE.print(f"[bold][yellow]:tada: Download failed due to {e}.")
+        raise e
+
+    return ckpt_path
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--repo_id", type=str, default=None, required=True)
+    parser.add_argument("--upload", action="store_true")
+    parser.add_argument("--download", action="store_true")
+    parser.add_argument("--folder_path", type=str, default=None, required=False)
+    parser.add_argument("--file_path", type=str, default=None, required=False)
+    parser.add_argument("--save_path", type=str, default=None, required=False)
+    parser.add_argument("--token", type=str, default=None, required=False)
+    args = parser.parse_args()
+
+    token = args.token or os.getenv("hf_token", None)
+    ignore_patterns = ["**/optimizer.bin", "**/random_states*", "**/scaler.pt", "**/scheduler.bin"]
+
+    if not (args.folder_path or args.file_path):
+        raise RuntimeError(f'Choose either folder path or file path please!')
+
+    if len(args.repo_id.split('/')) != 2:
+        raise RuntimeError(f'Invalid repo_id: {args.repo_id}, please use in [use-id]/[repo-name] format')
+    CONSOLE.log(f"Use repo: [bold][yellow] {args.repo_id}")
+
+    if args.upload:
+        upload()
+
+    if args.download:
+        download()

seq_10_150_3_3_encode_occ_separate_offsets_bs8_128_no_seqindex_long_lastvalid_ablation_grid/backups/scripts/pretrain_map.sh

@@ -0,0 +1,27 @@
+#!/bin/bash
+
+mkdir -p job_out
+
+#SBATCH --job-name YOUR_JOB_NAME        # Job name
+### Logging
+#SBATCH --output=job_out/%j.out                 # Stdout (%j expands to jobId)
+#SBATCH --error=job_out/%j.err                  # Stderr (%j expands to jobId)
+### Node info
+#SBATCH --nodes=1                       # Single node or multi node
+#SBATCH --nodelist=sota-6
+#SBATCH --time 20:00:00                 # Max time (hh:mm:ss)
+#SBATCH --gres=gpu:4                    # GPUs per node
+#SBATCH --mem=256G                       # Recommend 32G per GPU
+#SBATCH --ntasks-per-node=4             # Tasks per node
+#SBATCH --cpus-per-task=256               # Recommend 8 per GPU
+### Whatever your job needs to do
+
+export REQUESTS_CA_BUNDLE="/etc/ssl/certs/ca-certificates.crt"
+export HTTPS_PROXY="https://192.168.0.10:443/"
+export https_proxy="https://192.168.0.10:443/"
+
+export TEST_VAL_PRED=True
+export WANDB=True
+
+cd /u/xiuyu/work/dev4
+PYTHONPATH=".":$PYTHONPATH python3 train.py --config configs/train/pretrain_scalable_map.yaml --save_ckpt_path output/ours_map_pretrain

seq_10_150_3_3_encode_occ_separate_offsets_bs8_128_no_seqindex_long_lastvalid_ablation_grid/backups/scripts/run_eval.sh

@@ -0,0 +1,20 @@
+#! /bin/bash
+
+# env
+export REQUESTS_CA_BUNDLE="/etc/ssl/certs/ca-certificates.crt"
+export HTTPS_PROXY="https://192.168.0.10:443/"
+export https_proxy="https://192.168.0.10:443/"
+
+export WANDB=1
+
+# args
+DEVICES=$1
+CONFIG='configs/ours_long_term.yaml'
+# CKPT_PATH='output/scalable_smart_long/last.ckpt'
+CKPT_PATH='output2/seq_10_150_3_3_encode_occ_separate_offsets_bs8_128_no_seqindex_long/last.ckpt'
+
+# run
+PYTHONPATH=".":$PYTHONPATH python3 run.py \
+                            --devices $DEVICES \
+                            --config $CONFIG \
+                            --ckpt_path $CKPT_PATH ${@:2}

seq_10_150_3_3_encode_occ_separate_offsets_bs8_128_no_seqindex_long_lastvalid_ablation_grid/backups/scripts/run_train.sh

@@ -0,0 +1,20 @@
+#! /bin/bash
+
+# env
+export REQUESTS_CA_BUNDLE="/etc/ssl/certs/ca-certificates.crt"
+export HTTPS_PROXY="https://192.168.0.10:443/"
+export https_proxy="https://192.168.0.10:443/"
+
+export WANDB=1
+
+# args
+DEVICES=$1
+CONFIG='configs/ours_long_term.yaml'
+SAVE_CKPT_PATH='output/scalable_smart_long'
+
+# run
+PYTHONPATH=".":$PYTHONPATH python3 run.py \
+                            --train \
+                            --devices $DEVICES \
+                            --config $CONFIG \
+                            --save_ckpt_path $SAVE_CKPT_PATH