Set up self-contained Gradio Space

Bundle the app code and tiny demo dataset so the Hugging Face Space can boot directly into the SceneStreamer demo with sensible headless defaults.

Made-with: Cursor

.gitignore

@@ -0,0 +1,12 @@
+.venv/
+__pycache__/
+*.pyc
+*.pyo
+.DS_Store
+artifacts/
+outputs/
+lightning_logs/
+wandb/
+scenestreamer/outputs/
+scenestreamer/lightning_logs/
+scenestreamer/eval/outputs/

README.md

@@ -1,8 +1,8 @@
 ---
 title: SceneStreamer
-emoji: 😻
-colorFrom: green
-colorTo: gray
+emoji: 🚗
+colorFrom: blue
+colorTo: indigo
 sdk: gradio
 sdk_version: 6.9.0
 app_file: app.py
@@ -10,4 +10,28 @@ pinned: false
 license: mit
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# SceneStreamer Space
+
+This Space hosts the interactive Gradio demo for `SceneStreamer`.
+
+What is included here:
+
+- the Gradio app entrypoint in `app.py`
+- the SceneStreamer package code needed by the demo
+- a tiny bundled ScenarioNet subset in `data/20scenarios`
+
+Default behavior:
+
+- the app loads the bundled demo subset automatically
+- the model checkpoint is fetched from the Hugging Face Hub by default
+- `SCENESTREAMER_DEVICE` defaults to `cpu`
+
+Optional Space variables:
+
+- `SCENESTREAMER_DATASET_DIR`
+- `SCENESTREAMER_HF_REPO`
+- `SCENESTREAMER_HF_FILE`
+- `SCENESTREAMER_CKPT`
+- `SCENESTREAMER_DEVICE`
+
+If the app shows the setup screen instead of the demo, the dataset path is missing or the demo subset was not uploaded with the repo.

app.py

@@ -0,0 +1,58 @@
+from __future__ import annotations
+
+import os
+from pathlib import Path
+
+os.environ.setdefault("MPLBACKEND", "Agg")
+os.environ.setdefault("PYGAME_HIDE_SUPPORT_PROMPT", "1")
+os.environ.setdefault("SDL_AUDIODRIVER", "dummy")
+os.environ.setdefault("SDL_VIDEODRIVER", "dummy")
+
+import gradio as gr
+
+from scenestreamer.gradio_ui.demo_app import DEFAULT_HF_FILE, DEFAULT_HF_REPO, build_demo
+
+
+def _build_space_demo() -> gr.Blocks:
+    dataset_dir = os.environ.get("SCENESTREAMER_DATASET_DIR", "data/20scenarios")
+    hf_repo = os.environ.get("SCENESTREAMER_HF_REPO", DEFAULT_HF_REPO)
+    hf_file = os.environ.get("SCENESTREAMER_HF_FILE", DEFAULT_HF_FILE)
+    ckpt = os.environ.get("SCENESTREAMER_CKPT") or None
+    device = os.environ.get("SCENESTREAMER_DEVICE", "cpu")
+
+    if not Path(dataset_dir).exists():
+        with gr.Blocks(title="SceneStreamer Space Setup") as demo:
+            gr.Markdown("## SceneStreamer Space Setup Required")
+            gr.Markdown(
+                "This Space needs a local ScenarioNet dataset directory before the interactive demo can start.\n\n"
+                f"Current `SCENESTREAMER_DATASET_DIR`: `{dataset_dir}`"
+            )
+            gr.Markdown(
+                "Set Space variables or attach storage, then restart the Space:\n"
+                "- `SCENESTREAMER_DATASET_DIR`\n"
+                "- `SCENESTREAMER_HF_REPO` (optional)\n"
+                "- `SCENESTREAMER_HF_FILE` (optional)\n"
+                "- `SCENESTREAMER_CKPT` (optional local checkpoint)\n"
+                "- `SCENESTREAMER_DEVICE` (default `cpu`)"
+            )
+            gr.Markdown(
+                "This repo is expected to include a tiny bundled demo subset under `data/20scenarios`. "
+                "If you are seeing this page after pushing the repo, the demo data was likely not uploaded."
+            )
+        return demo
+
+    return build_demo(
+        dataset_dir=dataset_dir,
+        hf_repo=hf_repo,
+        hf_file=hf_file,
+        ckpt=ckpt,
+        device=device,
+    )
+
+
+demo = _build_space_demo()
+
+
+if __name__ == "__main__":
+    demo.launch()
+

cfgs/motion_default.yaml

@@ -0,0 +1,211 @@
+defaults:
+  - _self_
+
+# Experiment related
+exp_name: 'default'
+seed: 0
+epochs: 50
+batch_size: 10
+val_batch_size: 4
+num_workers: 16
+val_num_workers: 16
+num_sanity_val_steps: 100
+val_check_interval: 1.0
+wandb: False
+log_dir: Null
+limit_train_batches: -1
+limit_val_batches: -1
+prefetch_factor: 2
+ckpt: Null
+eval: False
+pretrain: Null
+deterministic: False
+detect_anomaly: False
+check_val_every_n_epoch: 1
+
+USE_RL_FINETUNING: False
+
+
+# Turn both on when training TG to match TrafficGen's metrics.
+LIMIT_MAP_RANGE: False
+FOLLOW_TRAFFICGEN: False
+FORCE_SDC_FOR_TRAFFICGEN: False
+ONLY_LANE_FOR_TRAFFICGEN: False
+
+# True then agent info will not be used in encoder,
+# and new tokens for history will be added for decoder.
+GPT_STYLE: false
+REMOVE_AGENT_FROM_SCENE_ENCODER: false
+USE_DIFFUSION: false
+USE_ADALN: null
+BACKWARD_PREDICTION: false
+
+
+USE_DESTINATION: false
+
+TF_DEST: True
+
+
+ADD_CONTOUR_RELATION: false
+
+DELTA_TOKENIZER_FILE_NAME: "1030_argsort_less_256_128_128.pkl"
+
+
+USE_TRAFFICGEN: false
+TRAIN_TRAFFICGEN: null
+USE_MOTION: true
+
+EVAL_MOTION: true
+EVAL_TRAFFICGEN: false
+
+DELTA_POS_IS_VELOCITY: false
+
+SIMPLE_RELATION: false
+SIMPLE_RELATION_FACTOR: 1
+
+RECONSTRUCT_MAP: false
+
+UPDATE_RELATION: false
+
+REMOVE_REL_NORM: false
+
+DATA:
+    TRAINING_DATA_DIR: '/data/datasets/scenarionet/waymo/training'
+    TEST_DATA_DIR: '/data/datasets/scenarionet/waymo/validation'
+    ADV_INFO_PATH: 'data/all_adv.pkl'
+    SAMPLE_INTERVAL_TRAINING: 1
+    SAMPLE_INTERVAL_TEST: 1
+    SD_PASSTHROUGH: false
+    ALLOW_CACHE: false
+    RETURN_HALFWAY: false  # Only used when generating LMDB dataset
+    USE_LMDB: false
+    USE_CACHE: false
+
+PREPROCESSING:
+    MAX_VECTORS: 128
+    MAX_MAP_FEATURES: 512
+    MAX_LENGTH_PER_MAP_FEATURE: 10000  # Useless
+    MAX_AGENTS: 128
+    MAX_TRAFFIC_LIGHTS: 64
+    PADDING_TO_MAX: false
+    keep_all_data: false # for debug
+    ADD_SDC_TO_OBJECT_OF_INTEREST: true  # Should be True when WOSAC
+    REMOVE_TRAFFIC_LIGHT_STATE: false
+    TRUNCATE_TIME: -1
+
+TRAINING:
+    PREDICT_ALL_AGENTS: true
+
+EVALUATION:
+    NAME: 'waymo_motion_prediction'
+    PREDICT_ALL_AGENTS: false
+    DELETE_EVAL_RESULT: true
+    NUM_MODES: 6
+    MAXIMUM_BATCH_SIZE: 10000
+    USE_CACHE: true
+    USE_TG_AS_GT: 1111
+    TG_REJECT_SAMPLING: True
+    TG_SDC_DISTANCE_MASKING: False
+
+MODEL:
+    NAME: 'motionlm'
+    D_MODEL: 256
+    NUM_ATTN_LAYERS: 4
+    NUM_ATTN_HEAD: 8
+#    DROPOUT_OF_ATTN: 0.0
+    DROPOUT: 0.0
+    NUM_DECODER_LAYERS: 6
+    ADD_PE_FOR_TOKEN: true
+    RELATIVE_PE: true
+    RELATIVE_PE_DECODER: false
+    PRE_PROJECTION: false
+    KNN: 128
+    S2S_DISTANCE: null
+    SELF_ATTN_KNN: 128
+    CROSS_ATTN_KNN: 128
+    RANDOMIZE_AGENT_ID: true
+    A2S_KNN: null
+    A2S_DISTANCE: null
+    A2A_KNN: null
+    A2A_DISTANCE: null
+    ADD_RELATION_TO_V: false
+    IS_V7: False
+    PER_CONTOUR_POINT_RELATION: null
+
+TOKENIZATION:
+    TOKENIZATION_METHOD: delta_delta
+    NUM_SKIPPED_STEPS: 5
+    NUM_BINS: 13
+    X_MAX: 3.5  # <<< Deprecated
+    X_MIN: -3.5  # <<< Deprecated
+    Y_MAX: 3.5  # <<< Deprecated
+    Y_MIN: -3.5  # <<< Deprecated
+    ADD_NOISE: false
+    NOISE_TOPK: 5
+    ALLOW_SKIP_STEP: True
+
+    MIN_DISPLACEMENT: 0.1
+    MIN_DISPLACEMENT_INIT: null
+    MIN_SPEED: null
+    SMOOTH_FACTOR: null
+    MAX_HEADING_DIFF: null
+
+    USE_CONTOUR_ERROR: True
+
+    VEH_LIMIT: 3.5
+    PED_LIMIT: 3.5
+    CYC_LIMIT: 3.5
+
+    FLIP_WRONG_HEADING: false
+    SHOULD_STANDARDIZE: true
+
+#    MIN_DISPLACEMENT: 0.3
+#    MIN_DISPLACEMENT_INIT: 1.0
+#    MIN_SPEED: 0.5
+#    SMOOTH_FACTOR: null
+#    MAX_HEADING_DIFF: 0.3
+
+
+SAMPLING:
+    SAMPLING_METHOD: 'topp'
+    TEMPERATURE: 1.0
+    TOPP: 0.95
+
+
+OPTIMIZATION:
+#    NUM_EPOCHS: 50
+    OPTIMIZER: AdamW
+    LR: 0.0003
+    WEIGHT_DECAY: 0.0
+    GRAD_NORM_CLIP: 1.0
+    SCHEDULER: cosine
+    WARMUP_STEPS: 2000
+#    TRAINING_STEPS: 300000
+    USE_FOCAL_LOSS: false
+
+
+SUBMISSION:
+    GENERATE_SUBMISSION: false
+    PREFIX: "peng"
+    ACCOUNT: "dr.zhenghao.peng@gmail.com"
+    METHOD_NAME: "peng"
+    num_model_parameters: '10m'  # TODO: Need to be changed accordingly!
+    SAVE_EVAL_DATA: true
+
+TMP_DIR: "tmp"  # Relative to repo root
+
+ACTION_LABEL:
+    USE_ACTION_LABEL: false  # Only valid for turning + acceleration
+    USE_SAFETY_LABEL: false
+    MASK_PROBABILITY_ACTION_LABEL: 0.0  # Might turn it on
+    MASK_PROBABILITY_SAFETY_LABEL: 0.0  # Might turn it on
+
+LANGUAGE_CONDITION: false
+FINE_TUNE_BERT: false
+
+MCTS:
+  USE_MCTS: False
+  MCTS_DEPTH: -1
+  MCTS_WIDTH: -1
+
+TOKEN_BUFFER_CACHE_LENGTH: 100

cfgs/scenestreamer-base-large.yaml

@@ -0,0 +1,96 @@
+defaults:
+  - motion_default
+  - _self_
+
+exp_name: 'scenestreamer-base-large'
+pretrain: "/bigdata/zhenghao/scenestreamer/lightning_logs/scenestreamer/20250506_scenestreamer_v18_notg_large_FIXETYPE_2025-05-06/checkpoints"
+
+num_workers: 8
+val_num_workers: 8
+num_sanity_val_steps: 10
+
+batch_size: 4
+val_batch_size: 4
+limit_val_batches: -1
+
+eval_backward_model: False
+
+epochs: 30
+wandb: True
+log_dir: /bigdata/zhenghao/scenestreamer
+
+SCENESTREAMER_ATTENTION_KNN: 128
+SCENESTREAMER_ATTENTION_MAX_DISTANCE: 50
+SCENESTREAMER_NO_TG: true
+
+REMOVE_AGENT_FROM_SCENE_ENCODER: True  # <<<
+
+BACKWARD_PREDICTION: False  # <<<
+ADD_CONTOUR_RELATION: True  # <<<
+
+DELTA_POS_IS_VELOCITY: True
+SIMPLE_RELATION: True
+
+RECONSTRUCT_MAP: False
+UPDATE_RELATION: False
+REMOVE_REL_NORM: False  # <<<
+
+USE_TRAFFICGEN: True
+USE_MOTION: True
+EVAL_MOTION: True
+EVAL_TRAFFICGEN: False
+
+GPT_STYLE: True  # <<<
+USE_ADALN: False
+
+SAMPLING:
+  TOPP: 0.95
+  TEMPERATURE: 1.0
+
+TOKENIZATION:
+  TOKENIZATION_METHOD: "BicycleModelTokenizerFixed0124"  # <<<
+  USE_CONTOUR_ERROR: True  # <<<
+  ALLOW_SKIP_STEP: True
+  ADD_NOISE: False
+  NUM_BINS: 33
+
+PREPROCESSING:
+    REMOVE_TRAFFIC_LIGHT_STATE: False
+    MAX_LENGTH_PER_MAP_FEATURE: 10
+    MAX_MAP_FEATURES: 3000
+    MAX_VECTORS: 30
+    MAX_AGENTS: 128
+    DEST_DROPOUT: 0.0
+    ADD_SDC_TO_OBJECT_OF_INTEREST: False
+
+DATA:
+  TRAINING_DATA_DIR: ''
+  TEST_DATA_DIR: ''
+
+MODEL:
+  USE_MOTION_HEAD_PRENORM: True
+  ALL_TO_MAP_3D: False
+  D_MODEL: 128
+  NAME: 'scenestreamer'
+  NUM_ATTN_HEAD: 4  # TODO: Need to increase? was 8
+  # Encoder:
+  NUM_ATTN_LAYERS: 2  # TODO: Need to increase? was 3
+  RELATIVE_PE: true
+  # Decoder:
+  NUM_DECODER_LAYERS: 4  # TODO: Need to increase? was 6
+  RELATIVE_PE_DECODER: True
+  SIMPLE_RELATION_FACTOR: 1
+  # New:
+  KNN: -100
+  S2S_DISTANCE: -100
+  A2S_KNN: -100
+  A2S_DISTANCE: -100
+  A2A_KNN: -100
+  A2A_DISTANCE: -100
+  ADD_RELATION_TO_V: False
+  PER_CONTOUR_POINT_RELATION: False
+  IS_V7: True
+
+SUBMISSION:
+  METHOD_NAME: "scenestreamer-base-large"
+  num_model_parameters: '3.3m'

cfgs/scenestreamer-base-small.yaml

@@ -0,0 +1,96 @@
+defaults:
+  - motion_default
+  - _self_
+
+exp_name: 'scenestreamer-base-small'
+pretrain: "/bigdata/zhenghao/scenestreamer/lightning_logs/scenestreamer/20250507_scenestreamer_v17_notg_finetune_FIXTYPE_2025-05-07/checkpoints"
+
+num_workers: 8
+val_num_workers: 8
+num_sanity_val_steps: 10
+
+batch_size: 4
+val_batch_size: 4
+limit_val_batches: -1
+
+eval_backward_model: False
+
+epochs: 30
+wandb: True
+log_dir: /bigdata/zhenghao/scenestreamer
+
+SCENESTREAMER_ATTENTION_KNN: 128
+SCENESTREAMER_ATTENTION_MAX_DISTANCE: 50
+SCENESTREAMER_NO_TG: true
+
+REMOVE_AGENT_FROM_SCENE_ENCODER: True  # <<<
+
+BACKWARD_PREDICTION: False  # <<<
+ADD_CONTOUR_RELATION: True  # <<<
+
+DELTA_POS_IS_VELOCITY: True
+SIMPLE_RELATION: True
+
+RECONSTRUCT_MAP: False
+UPDATE_RELATION: False
+REMOVE_REL_NORM: False  # <<<
+
+USE_TRAFFICGEN: True
+USE_MOTION: True
+EVAL_MOTION: True
+EVAL_TRAFFICGEN: False
+
+GPT_STYLE: True  # <<<
+USE_ADALN: False
+
+SAMPLING:
+  TOPP: 0.95
+  TEMPERATURE: 1.0
+
+TOKENIZATION:
+  TOKENIZATION_METHOD: "BicycleModelTokenizerFixed0124"  # <<<
+  USE_CONTOUR_ERROR: True  # <<<
+  ALLOW_SKIP_STEP: True
+  ADD_NOISE: False
+  NUM_BINS: 33
+
+PREPROCESSING:
+    REMOVE_TRAFFIC_LIGHT_STATE: False
+    MAX_LENGTH_PER_MAP_FEATURE: 10
+    MAX_MAP_FEATURES: 3000
+    MAX_VECTORS: 30
+    MAX_AGENTS: 128
+    DEST_DROPOUT: 0.0
+    ADD_SDC_TO_OBJECT_OF_INTEREST: False
+
+DATA:
+  TRAINING_DATA_DIR: ''
+  TEST_DATA_DIR: ''
+
+MODEL:
+  USE_MOTION_HEAD_PRENORM: True
+  ALL_TO_MAP_3D: False
+  D_MODEL: 64
+  NAME: 'scenestreamer'
+  NUM_ATTN_HEAD: 4  # TODO: Need to increase? was 8
+  # Encoder:
+  NUM_ATTN_LAYERS: 2  # TODO: Need to increase? was 3
+  RELATIVE_PE: true
+  # Decoder:
+  NUM_DECODER_LAYERS: 4  # TODO: Need to increase? was 6
+  RELATIVE_PE_DECODER: True
+  SIMPLE_RELATION_FACTOR: 1
+  # New:
+  KNN: -100
+  S2S_DISTANCE: -100
+  A2S_KNN: -100
+  A2S_DISTANCE: -100
+  A2A_KNN: -100
+  A2A_DISTANCE: -100
+  ADD_RELATION_TO_V: False
+  PER_CONTOUR_POINT_RELATION: False
+  IS_V7: True
+
+SUBMISSION:
+  METHOD_NAME: "scenestreamer-base-small"
+  num_model_parameters: '1.1m'

cfgs/scenestreamer-base-xl.yaml

@@ -0,0 +1,97 @@
+defaults:
+  - motion_default
+  - _self_
+
+exp_name: 'scenestreamer-base-xl'
+pretrain: "/bigdata/zhenghao/scenestreamer/lightning_logs/scenestreamer/20250512_scenestreamer-base-xl_2025-05-12/checkpoints"
+
+num_workers: 8
+val_num_workers: 8
+num_sanity_val_steps: 10
+
+batch_size: 4
+val_batch_size: 4
+limit_val_batches: -1
+
+eval_backward_model: False
+
+epochs: 30
+wandb: True
+log_dir: /bigdata/zhenghao/scenestreamer
+
+SCENESTREAMER_ATTENTION_KNN: 128
+SCENESTREAMER_ATTENTION_MAX_DISTANCE: 50
+SCENESTREAMER_NO_TG: true
+
+REMOVE_AGENT_FROM_SCENE_ENCODER: True  # <<<
+
+BACKWARD_PREDICTION: False  # <<<
+ADD_CONTOUR_RELATION: True  # <<<
+
+DELTA_POS_IS_VELOCITY: True
+SIMPLE_RELATION: True
+
+RECONSTRUCT_MAP: False
+UPDATE_RELATION: False
+REMOVE_REL_NORM: False  # <<<
+
+USE_TRAFFICGEN: True
+USE_MOTION: True
+EVAL_MOTION: True
+EVAL_TRAFFICGEN: False
+
+GPT_STYLE: True  # <<<
+USE_ADALN: False
+
+SAMPLING:
+  TOPP: 0.95
+  TEMPERATURE: 1.0
+
+TOKENIZATION:
+  TOKENIZATION_METHOD: "BicycleModelTokenizerFixed0124"  # <<<
+  USE_CONTOUR_ERROR: True  # <<<
+  ALLOW_SKIP_STEP: True
+  ADD_NOISE: False
+  NUM_BINS: 33
+
+PREPROCESSING:
+  REMOVE_TRAFFIC_LIGHT_STATE: False
+  MAX_LENGTH_PER_MAP_FEATURE: 10
+  MAX_MAP_FEATURES: 3000
+  MAX_VECTORS: 30
+  MAX_AGENTS: 128
+  DEST_DROPOUT: 0.0
+  ADD_SDC_TO_OBJECT_OF_INTEREST: False
+
+DATA:
+  TRAINING_DATA_DIR: ''
+  TEST_DATA_DIR: ''
+
+MODEL:
+  USE_MOTION_HEAD_PRENORM: True
+  ALL_TO_MAP_3D: False
+  D_MODEL: 128
+  NAME: 'scenestreamer'
+  NUM_ATTN_HEAD: 8
+  # Encoder:
+  NUM_ATTN_LAYERS: 3
+  RELATIVE_PE: true
+  # Decoder:
+  NUM_DECODER_LAYERS: 6
+  RELATIVE_PE_DECODER: True
+  SIMPLE_RELATION_FACTOR: 1
+  # New:
+  KNN: -100
+  S2S_DISTANCE: -100
+  A2S_KNN: -100
+  A2S_DISTANCE: -100
+  A2A_KNN: -100
+  A2A_DISTANCE: -100
+  ADD_RELATION_TO_V: False
+  PER_CONTOUR_POINT_RELATION: False
+  IS_V7: True
+
+SUBMISSION:
+  METHOD_NAME: "scenestreamer-base-xl"
+  num_model_parameters: '4.2m'
+  ACCOUNT: "dr.zhenghao.peng@gmail.com"

cfgs/scenestreamer-full-large-nors.yaml

@@ -0,0 +1,99 @@
+defaults:
+  - motion_default
+  - _self_
+
+exp_name: 'scenestreamer-full-large-nors'
+pretrain: "/bigdata/zhenghao/scenestreamer/lightning_logs/scenestreamer/20250507_scenestreamer-full-large_2025-05-07/checkpoints/20250507_scenestreamer-full-large_2025-05-07_epoch=1-step=77031.ckpt"
+
+num_workers: 8
+val_num_workers: 8
+num_sanity_val_steps: 10
+
+batch_size: 4
+val_batch_size: 4
+limit_val_batches: -1
+
+eval_backward_model: False
+
+epochs: 30
+wandb: True
+log_dir: /bigdata/zhenghao/scenestreamer
+
+SCENESTREAMER_ATTENTION_KNN: 128
+SCENESTREAMER_ATTENTION_MAX_DISTANCE: 50
+SCENESTREAMER_NO_TG: false
+
+REMOVE_AGENT_FROM_SCENE_ENCODER: True  # <<<
+
+BACKWARD_PREDICTION: False  # <<<
+ADD_CONTOUR_RELATION: True  # <<<
+
+DELTA_POS_IS_VELOCITY: True
+SIMPLE_RELATION: True
+
+RECONSTRUCT_MAP: False
+UPDATE_RELATION: False
+REMOVE_REL_NORM: False  # <<<
+
+USE_TRAFFICGEN: True
+USE_MOTION: True
+EVAL_MOTION: True
+EVAL_TRAFFICGEN: False
+
+GPT_STYLE: True  # <<<
+USE_ADALN: False
+
+SAMPLING:
+  TOPP: 0.95
+  TEMPERATURE: 1.0
+
+TOKENIZATION:
+  TOKENIZATION_METHOD: "BicycleModelTokenizerFixed0124"  # <<<
+  USE_CONTOUR_ERROR: True  # <<<
+  ALLOW_SKIP_STEP: True
+  ADD_NOISE: False
+  NUM_BINS: 33
+
+PREPROCESSING:
+  REMOVE_TRAFFIC_LIGHT_STATE: False
+  MAX_LENGTH_PER_MAP_FEATURE: 10
+  MAX_MAP_FEATURES: 3000
+  MAX_VECTORS: 30
+  MAX_AGENTS: 128
+  DEST_DROPOUT: 0.0
+  ADD_SDC_TO_OBJECT_OF_INTEREST: False
+
+DATA:
+  TRAINING_DATA_DIR: ''
+  TEST_DATA_DIR: ''
+
+MODEL:
+  USE_MOTION_HEAD_PRENORM: True
+  ALL_TO_MAP_3D: False
+  D_MODEL: 128
+  NAME: 'scenestreamer'
+  NUM_ATTN_HEAD: 4  # TODO: Need to increase? was 8
+  # Encoder:
+  NUM_ATTN_LAYERS: 2  # TODO: Need to increase? was 3
+  RELATIVE_PE: true
+  # Decoder:
+  NUM_DECODER_LAYERS: 4  # TODO: Need to increase? was 6
+  RELATIVE_PE_DECODER: True
+  SIMPLE_RELATION_FACTOR: 1
+  # New:
+  KNN: -100
+  S2S_DISTANCE: -100
+  A2S_KNN: -100
+  A2S_DISTANCE: -100
+  A2A_KNN: -100
+  A2A_DISTANCE: -100
+  ADD_RELATION_TO_V: False
+  PER_CONTOUR_POINT_RELATION: False
+  IS_V7: True
+
+SUBMISSION:
+  METHOD_NAME: "scenestreamer-full-large"
+  num_model_parameters: '4.6m'
+
+EVALUATION:
+  TG_REJECT_SAMPLING: False

cfgs/scenestreamer-full-large.yaml

@@ -0,0 +1,99 @@
+defaults:
+  - motion_default
+  - _self_
+
+exp_name: 'scenestreamer-full-large'
+pretrain: "/bigdata/zhenghao/scenestreamer/lightning_logs/scenestreamer/20250507_scenestreamer-full-large_2025-05-07/checkpoints/20250507_scenestreamer-full-large_2025-05-07_epoch=1-step=77031.ckpt"
+
+num_workers: 8
+val_num_workers: 8
+num_sanity_val_steps: 10
+
+batch_size: 4
+val_batch_size: 4
+limit_val_batches: -1
+
+eval_backward_model: False
+
+epochs: 30
+wandb: True
+log_dir: /bigdata/zhenghao/scenestreamer
+
+SCENESTREAMER_ATTENTION_KNN: 128
+SCENESTREAMER_ATTENTION_MAX_DISTANCE: 50
+SCENESTREAMER_NO_TG: false
+
+REMOVE_AGENT_FROM_SCENE_ENCODER: True  # <<<
+
+BACKWARD_PREDICTION: False  # <<<
+ADD_CONTOUR_RELATION: True  # <<<
+
+DELTA_POS_IS_VELOCITY: True
+SIMPLE_RELATION: True
+
+RECONSTRUCT_MAP: False
+UPDATE_RELATION: False
+REMOVE_REL_NORM: False  # <<<
+
+USE_TRAFFICGEN: True
+USE_MOTION: True
+EVAL_MOTION: True
+EVAL_TRAFFICGEN: False
+
+GPT_STYLE: True  # <<<
+USE_ADALN: False
+
+SAMPLING:
+  TOPP: 0.95
+  TEMPERATURE: 1.0
+
+TOKENIZATION:
+  TOKENIZATION_METHOD: "BicycleModelTokenizerFixed0124"  # <<<
+  USE_CONTOUR_ERROR: True  # <<<
+  ALLOW_SKIP_STEP: True
+  ADD_NOISE: False
+  NUM_BINS: 33
+
+PREPROCESSING:
+    REMOVE_TRAFFIC_LIGHT_STATE: False
+    MAX_LENGTH_PER_MAP_FEATURE: 10
+    MAX_MAP_FEATURES: 3000
+    MAX_VECTORS: 30
+    MAX_AGENTS: 128
+    DEST_DROPOUT: 0.0
+    ADD_SDC_TO_OBJECT_OF_INTEREST: False
+
+DATA:
+  TRAINING_DATA_DIR: ''
+  TEST_DATA_DIR: ''
+
+MODEL:
+  USE_MOTION_HEAD_PRENORM: True
+  ALL_TO_MAP_3D: False
+  D_MODEL: 128
+  NAME: 'scenestreamer'
+  NUM_ATTN_HEAD: 4  # TODO: Need to increase? was 8
+  # Encoder:
+  NUM_ATTN_LAYERS: 2  # TODO: Need to increase? was 3
+  RELATIVE_PE: true
+  # Decoder:
+  NUM_DECODER_LAYERS: 4  # TODO: Need to increase? was 6
+  RELATIVE_PE_DECODER: True
+  SIMPLE_RELATION_FACTOR: 1
+  # New:
+  KNN: -100
+  S2S_DISTANCE: -100
+  A2S_KNN: -100
+  A2S_DISTANCE: -100
+  A2A_KNN: -100
+  A2A_DISTANCE: -100
+  ADD_RELATION_TO_V: False
+  PER_CONTOUR_POINT_RELATION: False
+  IS_V7: True
+
+SUBMISSION:
+  METHOD_NAME: "scenestreamer-full-large"
+  num_model_parameters: '4.6m'
+
+EVALUATION:
+  TG_REJECT_SAMPLING: True

cfgs/scenestreamer-full-small.yaml

@@ -0,0 +1,96 @@
+defaults:
+  - motion_default
+  - _self_
+
+exp_name: 'scenestreamer-full-small'
+pretrain: "/bigdata/zhenghao/scenestreamer/lightning_logs/scenestreamer/20250505_scenestreamer_v19_withtg_nodest_FIXEDAS_2025-05-05/checkpoints"
+
+num_workers: 8
+val_num_workers: 8
+num_sanity_val_steps: 10
+
+batch_size: 4
+val_batch_size: 4
+limit_val_batches: -1
+
+eval_backward_model: False
+
+epochs: 30
+wandb: True
+log_dir: /bigdata/zhenghao/scenestreamer
+
+SCENESTREAMER_ATTENTION_KNN: 128
+SCENESTREAMER_ATTENTION_MAX_DISTANCE: 50
+SCENESTREAMER_NO_TG: false
+
+REMOVE_AGENT_FROM_SCENE_ENCODER: True  # <<<
+
+BACKWARD_PREDICTION: False  # <<<
+ADD_CONTOUR_RELATION: True  # <<<
+
+DELTA_POS_IS_VELOCITY: True
+SIMPLE_RELATION: True
+
+RECONSTRUCT_MAP: False
+UPDATE_RELATION: False
+REMOVE_REL_NORM: False  # <<<
+
+USE_TRAFFICGEN: True
+USE_MOTION: True
+EVAL_MOTION: True
+EVAL_TRAFFICGEN: False
+
+GPT_STYLE: True  # <<<
+USE_ADALN: False
+
+SAMPLING:
+  TOPP: 0.95
+  TEMPERATURE: 1.0
+
+TOKENIZATION:
+  TOKENIZATION_METHOD: "BicycleModelTokenizerFixed0124"  # <<<
+  USE_CONTOUR_ERROR: True  # <<<
+  ALLOW_SKIP_STEP: True
+  ADD_NOISE: False
+  NUM_BINS: 33
+
+PREPROCESSING:
+    REMOVE_TRAFFIC_LIGHT_STATE: False
+    MAX_LENGTH_PER_MAP_FEATURE: 10
+    MAX_MAP_FEATURES: 3000
+    MAX_VECTORS: 30
+    MAX_AGENTS: 128
+    DEST_DROPOUT: 0.0
+    ADD_SDC_TO_OBJECT_OF_INTEREST: False
+
+DATA:
+  TRAINING_DATA_DIR: ''
+  TEST_DATA_DIR: ''
+
+MODEL:
+  USE_MOTION_HEAD_PRENORM: True
+  ALL_TO_MAP_3D: False
+  D_MODEL: 64  # TODO: Need to increase? was 128
+  NAME: 'scenestreamer'
+  NUM_ATTN_HEAD: 4  # TODO: Need to increase? was 8
+  # Encoder:
+  NUM_ATTN_LAYERS: 2  # TODO: Need to increase? was 3
+  RELATIVE_PE: true
+  # Decoder:
+  NUM_DECODER_LAYERS: 4  # TODO: Need to increase? was 6
+  RELATIVE_PE_DECODER: True
+  SIMPLE_RELATION_FACTOR: 1
+  # New:
+  KNN: -100
+  S2S_DISTANCE: -100
+  A2S_KNN: -100
+  A2S_DISTANCE: -100
+  A2A_KNN: -100
+  A2A_DISTANCE: -100
+  ADD_RELATION_TO_V: False
+  PER_CONTOUR_POINT_RELATION: False
+  IS_V7: True
+
+SUBMISSION:
+  METHOD_NAME: "scenestreamer-full-small"
+  num_model_parameters: '1.5m'

cfgs/scenestreamer-full-xl.yaml

@@ -0,0 +1,100 @@
+defaults:
+  - motion_default
+  - _self_
+
+exp_name: 'scenestreamer-full-xl'
+pretrain: "/bigdata/zhenghao/scenestreamer/lightning_logs/scenestreamer/20250518_scenestreamer-full-xs_2025-05-18/checkpoints/20250518_scenestreamer-full-xs_2025-05-18_epoch=1-step=62494.ckpt"
+
+num_workers: 8
+val_num_workers: 8
+num_sanity_val_steps: 10
+
+batch_size: 4
+val_batch_size: 4
+limit_val_batches: -1
+
+eval_backward_model: False
+
+epochs: 30
+wandb: True
+log_dir: /bigdata/zhenghao/scenestreamer
+
+SCENESTREAMER_ATTENTION_KNN: 128
+SCENESTREAMER_ATTENTION_MAX_DISTANCE: 50
+SCENESTREAMER_NO_TG: false
+
+REMOVE_AGENT_FROM_SCENE_ENCODER: True  # <<<
+
+BACKWARD_PREDICTION: False  # <<<
+ADD_CONTOUR_RELATION: True  # <<<
+
+DELTA_POS_IS_VELOCITY: True
+SIMPLE_RELATION: True
+
+RECONSTRUCT_MAP: False
+UPDATE_RELATION: False
+REMOVE_REL_NORM: False  # <<<
+
+USE_TRAFFICGEN: True
+USE_MOTION: True
+EVAL_MOTION: True
+EVAL_TRAFFICGEN: False
+
+GPT_STYLE: True  # <<<
+USE_ADALN: False
+
+SAMPLING:
+  TOPP: 0.95
+  TEMPERATURE: 1.0
+
+TOKENIZATION:
+  TOKENIZATION_METHOD: "BicycleModelTokenizerFixed0124"  # <<<
+  USE_CONTOUR_ERROR: True  # <<<
+  ALLOW_SKIP_STEP: True
+  ADD_NOISE: False
+  NUM_BINS: 33
+
+PREPROCESSING:
+  REMOVE_TRAFFIC_LIGHT_STATE: False
+  MAX_LENGTH_PER_MAP_FEATURE: 10
+  MAX_MAP_FEATURES: 3000
+  MAX_VECTORS: 30
+  MAX_AGENTS: 128
+  DEST_DROPOUT: 0.0
+  ADD_SDC_TO_OBJECT_OF_INTEREST: False
+
+DATA:
+  TRAINING_DATA_DIR: ''
+  TEST_DATA_DIR: ''
+
+MODEL:
+  USE_MOTION_HEAD_PRENORM: True
+  ALL_TO_MAP_3D: False
+  D_MODEL: 128
+  NAME: 'scenestreamer'
+  NUM_ATTN_HEAD: 8
+  # Encoder:
+  NUM_ATTN_LAYERS: 3
+  RELATIVE_PE: true
+  # Decoder:
+  NUM_DECODER_LAYERS: 6
+  RELATIVE_PE_DECODER: True
+  SIMPLE_RELATION_FACTOR: 1
+  # New:
+  KNN: -100
+  S2S_DISTANCE: -100
+  A2S_KNN: -100
+  A2S_DISTANCE: -100
+  A2A_KNN: -100
+  A2A_DISTANCE: -100
+  ADD_RELATION_TO_V: False
+  PER_CONTOUR_POINT_RELATION: False
+  IS_V7: True
+
+SUBMISSION:
+  METHOD_NAME: "scenestreamer-full-xl"
+  num_model_parameters: '5.5m'
+  ACCOUNT: "dr.zhenghao.peng@gmail.com"
+
+EVALUATION:
+  TG_REJECT_SAMPLING: True

data/20scenarios/cache/sd_training.tfrecord-00000-of-01000_101d4e5775093d0c.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d2f5a926ba159d4e9acec464c9d091c093d138a21796ce5c264fea7f4398a777
+size 3007314

data/20scenarios/cache/sd_training.tfrecord-00000-of-01000_16d4f837fb57ff1.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b8a001efb1d464a0f9a3c06f26cf921f2308e81b26a65741491875201ede70b1
+size 6364095

data/20scenarios/cache/sd_training.tfrecord-00000-of-01000_17100879b93ceb61.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:742aa8c350793d83e949396dcb055a17bcbdd4b2b728b41d5f1b5c6c5a897ce1
+size 4382994

data/20scenarios/cache/sd_training.tfrecord-00000-of-01000_c403d5992cab9e0.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:716fa8212d8d4dbed60703cf5a9a952129fab5ed9342da748a7a48f00478e6d9
+size 11279523

data/20scenarios/process.ipynb

@@ -0,0 +1,128 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "id": "f003e6e4",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import os, pickle"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "id": "fdb6f299",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\u001b[0m\u001b[01;32mdataset_summary.pkl\u001b[0m*\r\n",
+      "\u001b[01;32mprocess.ipynb\u001b[0m*\r\n",
+      "\u001b[01;32msd_training.tfrecord-00000-of-01000_101d4e5775093d0c.pkl\u001b[0m*\r\n",
+      "\u001b[01;32msd_training.tfrecord-00000-of-01000_10c3969f1eb158d.pkl\u001b[0m*\r\n",
+      "\u001b[01;32msd_training.tfrecord-00000-of-01000_1109b0038ed8f25a.pkl\u001b[0m*\r\n",
+      "\u001b[01;32msd_training.tfrecord-00000-of-01000_116d257e98878d94.pkl\u001b[0m*\r\n",
+      "\u001b[01;32msd_training.tfrecord-00000-of-01000_11cdace2c1445900.pkl\u001b[0m*\r\n",
+      "\u001b[01;32msd_training.tfrecord-00000-of-01000_12a40f114f1ec5fc.pkl\u001b[0m*\r\n",
+      "\u001b[01;32msd_training.tfrecord-00000-of-01000_149f682e19454efa.pkl\u001b[0m*\r\n",
+      "\u001b[01;32msd_training.tfrecord-00000-of-01000_16361c8c522cf0e.pkl\u001b[0m*\r\n",
+      "\u001b[01;32msd_training.tfrecord-00000-of-01000_1663adf01133d82.pkl\u001b[0m*\r\n",
+      "\u001b[01;32msd_training.tfrecord-00000-of-01000_16b7c9a8ae6e89f1.pkl\u001b[0m*\r\n",
+      "\u001b[01;32msd_training.tfrecord-00000-of-01000_16d4f837fb57ff1.pkl\u001b[0m*\r\n",
+      "\u001b[01;32msd_training.tfrecord-00000-of-01000_17100879b93ceb61.pkl\u001b[0m*\r\n",
+      "\u001b[01;32msd_training.tfrecord-00000-of-01000_17fa718cd1251a26.pkl\u001b[0m*\r\n",
+      "\u001b[01;32msd_training.tfrecord-00000-of-01000_18840a098288507f.pkl\u001b[0m*\r\n",
+      "\u001b[01;32msd_training.tfrecord-00000-of-01000_18f5ce249ee2e949.pkl\u001b[0m*\r\n",
+      "\u001b[01;32msd_training.tfrecord-00000-of-01000_1a7143a44e480ca6.pkl\u001b[0m*\r\n",
+      "\u001b[01;32msd_training.tfrecord-00000-of-01000_1a8cc570d620bd31.pkl\u001b[0m*\r\n",
+      "\u001b[01;32msd_training.tfrecord-00000-of-01000_2a1e44d405a6833f.pkl\u001b[0m*\r\n",
+      "\u001b[01;32msd_training.tfrecord-00000-of-01000_8a346109094cd5aa.pkl\u001b[0m*\r\n",
+      "\u001b[01;32msd_training.tfrecord-00000-of-01000_c403d5992cab9e0.pkl\u001b[0m*\r\n"
+     ]
+    }
+   ],
+   "source": [
+    "ls"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "id": "f08caf3b",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "['sd_training.tfrecord-00000-of-01000_101d4e5775093d0c.pkl',\n",
+       " 'sd_training.tfrecord-00000-of-01000_10c3969f1eb158d.pkl',\n",
+       " 'sd_training.tfrecord-00000-of-01000_1109b0038ed8f25a.pkl',\n",
+       " 'sd_training.tfrecord-00000-of-01000_116d257e98878d94.pkl',\n",
+       " 'sd_training.tfrecord-00000-of-01000_11cdace2c1445900.pkl',\n",
+       " 'sd_training.tfrecord-00000-of-01000_12a40f114f1ec5fc.pkl',\n",
+       " 'sd_training.tfrecord-00000-of-01000_149f682e19454efa.pkl',\n",
+       " 'sd_training.tfrecord-00000-of-01000_16361c8c522cf0e.pkl',\n",
+       " 'sd_training.tfrecord-00000-of-01000_1663adf01133d82.pkl',\n",
+       " 'sd_training.tfrecord-00000-of-01000_16b7c9a8ae6e89f1.pkl',\n",
+       " 'sd_training.tfrecord-00000-of-01000_16d4f837fb57ff1.pkl',\n",
+       " 'sd_training.tfrecord-00000-of-01000_17100879b93ceb61.pkl',\n",
+       " 'sd_training.tfrecord-00000-of-01000_17fa718cd1251a26.pkl',\n",
+       " 'sd_training.tfrecord-00000-of-01000_18840a098288507f.pkl',\n",
+       " 'sd_training.tfrecord-00000-of-01000_18f5ce249ee2e949.pkl',\n",
+       " 'sd_training.tfrecord-00000-of-01000_1a7143a44e480ca6.pkl',\n",
+       " 'sd_training.tfrecord-00000-of-01000_1a8cc570d620bd31.pkl',\n",
+       " 'sd_training.tfrecord-00000-of-01000_2a1e44d405a6833f.pkl',\n",
+       " 'sd_training.tfrecord-00000-of-01000_8a346109094cd5aa.pkl',\n",
+       " 'sd_training.tfrecord-00000-of-01000_c403d5992cab9e0.pkl']"
+      ]
+     },
+     "execution_count": 3,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "[p for p in os.listdir(\".\") if p.endswith(\".pkl\") and p.startswith(\"sd\")]"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "id": "d1d50b21",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "d = {}\n",
+    "for p in [p for p in os.listdir(\".\") if p.endswith(\".pkl\") and p.startswith(\"sd\")]:\n",
+    "    d[p] = {}\n",
+    "\n",
+    "pickle.dump(d, open(\"dataset_summary.pkl\", \"wb\"))"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3 (ipykernel)",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.9.16"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

data/20scenarios/sd_training.tfrecord-00000-of-01000_101d4e5775093d0c.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:de3ceb564e944fda1d5b8e2f72428ada2b790d5d57195e6d9bca2cdf761a37f0
+size 338026

data/20scenarios/sd_training.tfrecord-00000-of-01000_10c3969f1eb158d.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:72cae539b2993d4dd668e5029460e2cdbf621dc1734d69ce30b357a156bc8375
+size 535617

data/20scenarios/sd_training.tfrecord-00000-of-01000_1109b0038ed8f25a.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:383b8cdeaefc13c15751f3ff29d28426543c24b6de674e0c5434e39f5d7a8d1f
+size 1060327

data/20scenarios/sd_training.tfrecord-00000-of-01000_116d257e98878d94.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:77209cb776b340b849ff123c7c9019d17495df41bf69fe2cfc765d5c8235fc66
+size 680549

data/20scenarios/sd_training.tfrecord-00000-of-01000_11cdace2c1445900.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:bf6fc9426a5b209d4f0862ff8046844f40fbf822fcf54a24947f23aa0f140161
+size 426242

data/20scenarios/sd_training.tfrecord-00000-of-01000_12a40f114f1ec5fc.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ddfd0775eedec97b30316beb8c537e3d768928abbd255030403c5bdff59a7f85
+size 837708

data/20scenarios/sd_training.tfrecord-00000-of-01000_149f682e19454efa.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3b2f40bfe73b50954987f15a462a4a86659615e416b577b1c1096c8c69ef0d05
+size 667369

data/20scenarios/sd_training.tfrecord-00000-of-01000_16361c8c522cf0e.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:30f1f5e1d3a5863a88b919a79c8c4a75bff2774df99412a092706d5cebc7ffbb
+size 319336

data/20scenarios/sd_training.tfrecord-00000-of-01000_1663adf01133d82.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:bacb875cdcd59d4c5ada7c1a2b546a31967de6f15f288d81136d2f3cc12f2413
+size 560466

data/20scenarios/sd_training.tfrecord-00000-of-01000_16b7c9a8ae6e89f1.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5194b6697b1fa3d00113c8ba2f39d50ffb8066b004c8f2f5d15953a550027e7c
+size 456308

data/20scenarios/sd_training.tfrecord-00000-of-01000_16d4f837fb57ff1.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:dfce54f80da166a8cbaf5fe167770975dd4dfcc9bba3c29ad228ad112d995690
+size 396202

data/20scenarios/sd_training.tfrecord-00000-of-01000_17100879b93ceb61.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a15fb32533b7472c38ec32c676a9425490de2c95d8dd39b9aa2a6b98f1050512
+size 253625

data/20scenarios/sd_training.tfrecord-00000-of-01000_17fa718cd1251a26.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d126e0c7c1e4afa561739437ad7c7810830c766a45b725fb8cb8a33015029e61
+size 490897

data/20scenarios/sd_training.tfrecord-00000-of-01000_18840a098288507f.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ffbd89555db615f2d20516fedf2dd3689b4487221a82c9f48251e29c30027f31
+size 627272

data/20scenarios/sd_training.tfrecord-00000-of-01000_18f5ce249ee2e949.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4e93b1b9cbfd8c541c5c52523aee09f5abd3410d217bf6e402ec217935f5df08
+size 435189

data/20scenarios/sd_training.tfrecord-00000-of-01000_1a7143a44e480ca6.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:11ac7db4a7c917c8045c957ee93d2f761de089152fdb9b03554a889d4f1deca3
+size 376448

data/20scenarios/sd_training.tfrecord-00000-of-01000_1a8cc570d620bd31.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5cb753a3c0d7255b5d4333a95c5555b4b159dcc600f9da286787ab37893c70d8
+size 1558365

data/20scenarios/sd_training.tfrecord-00000-of-01000_2a1e44d405a6833f.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7ac83f07304e39a8c5e468399b12e3ed034d6855b2921ab094ee425e492bea6d
+size 1319602

data/20scenarios/sd_training.tfrecord-00000-of-01000_8a346109094cd5aa.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c81210cec08ec7bc5813e59bf4a2f8287800c2d626929441a0e7398f7283c509
+size 1316507

data/20scenarios/sd_training.tfrecord-00000-of-01000_c403d5992cab9e0.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:26feb91313bee0c89ff1a9660fc352ea687daf04a7ca642f600d2ece17a6e301
+size 1079778

packages.txt

@@ -0,0 +1,7 @@
+ffmpeg
+libgl1
+libglib2.0-0
+libsm6
+libxext6
+libxrender1
+libsdl2-2.0-0

pyproject.toml

@@ -0,0 +1,80 @@
+[project]
+name = "scenestreamer"
+version = "1.0.0"
+description = "SceneStreamer: Continuous Scenario Generation as Next Token Group Prediction"
+readme = "README.md"
+license = {text = "MIT"}
+authors = [
+    {name = "Zhenghao Peng", email = "pzh@berkeley.edu"}
+]
+requires-python = ">=3.10,<3.12"
+classifiers = [
+    "Development Status :: 4 - Beta",
+    "License :: OSI Approved :: MIT License",
+    "Programming Language :: Python :: 3.10",
+    "Programming Language :: Python :: 3.11",
+    "Topic :: Scientific/Engineering :: Artificial Intelligence",
+]
+
+dependencies = [
+    "torch>=2.0.0",
+    "torchvision",
+    "lightning>=2.0.0",
+    "hydra-core",
+    "omegaconf",
+    "numpy",
+    "tqdm",
+    "matplotlib",
+    "seaborn",
+    "Pillow",
+    "easydict",
+    "wandb",
+    "torch_geometric",
+    "transformers",
+    "tokenizers",
+    "huggingface_hub",
+    "tensorboardX",
+    "pyyaml",
+    "scikit-image",
+    "chardet",
+    "charset-normalizer",
+    "tabulate",
+    "metadrive-simulator",
+    "gradio>=6.9.0",
+    "scenarionet @ git+https://github.com/metadriverse/scenarionet.git",
+]
+
+[project.scripts]
+scenestreamer = "scenestreamer.cli:main"
+
+[project.optional-dependencies]
+dev = [
+    "ruff",
+    "pytest",
+]
+rl = [
+    "stable-baselines3>=2.0.0",
+    "gymnasium>=0.29.0",
+    "ipython",
+]
+# Note: waymo-open-dataset requires Python 3.10 and specific numpy versions.
+# Install separately: pip install waymo-open-dataset-tf-2-12-0==1.6.4
+
+[project.urls]
+Homepage = "https://vail-ucla.github.io/scenestreamer/"
+Repository = "https://github.com/pengzhenghao/scenestreamer"
+
+[build-system]
+requires = ["setuptools>=61.0"]
+build-backend = "setuptools.build_meta"
+
+[tool.setuptools.packages.find]
+include = ["scenestreamer*"]
+
+[tool.ruff]
+line-length = 120
+target-version = "py310"
+
+[tool.ruff.lint]
+select = ["E", "F", "W"]
+ignore = ["E501"]

requirements.txt

@@ -0,0 +1 @@
+.

scenestreamer/cli.py

@@ -0,0 +1,293 @@
+from __future__ import annotations
+
+import argparse
+import json
+import os
+import pathlib
+import runpy
+import sys
+import time
+from dataclasses import dataclass
+from typing import Any
+
+import yaml
+
+
+def _to_plain(obj: Any) -> Any:
+    if hasattr(obj, "items"):
+        return {k: _to_plain(v) for k, v in obj.items()}
+    if isinstance(obj, (list, tuple)):
+        return [_to_plain(v) for v in obj]
+    return obj
+
+
+def _to_easydict(obj: Any):
+    from easydict import EasyDict
+
+    if isinstance(obj, dict):
+        return EasyDict({k: _to_easydict(v) for k, v in obj.items()})
+    if isinstance(obj, list):
+        return [_to_easydict(v) for v in obj]
+    return obj
+
+
+def load_yaml_config(path: str | os.PathLike[str]):
+    with open(path, "r") as f:
+        data = yaml.safe_load(f)
+    return _to_easydict(data)
+
+
+def apply_overrides(cfg, overrides: list[str]) -> None:
+    """
+    Apply overrides of form KEY=VALUE where KEY is dot-delimited.
+    VALUE is parsed using yaml.safe_load (so numbers/bools/lists work).
+    """
+    for item in overrides:
+        if "=" not in item:
+            raise ValueError(f"Invalid override (expected KEY=VALUE): {item}")
+        key, raw_val = item.split("=", 1)
+        value = yaml.safe_load(raw_val)
+
+        cur = cfg
+        parts = key.split(".")
+        for p in parts[:-1]:
+            if not hasattr(cur, p):
+                setattr(cur, p, _to_easydict({}))
+            cur = getattr(cur, p)
+        setattr(cur, parts[-1], value)
+
+
+@dataclass(frozen=True)
+class RunPaths:
+    run_dir: pathlib.Path
+    config_path: pathlib.Path
+    metrics_path: pathlib.Path
+
+
+def make_run_dir(base_dir: str | os.PathLike[str], run_id: str | None) -> RunPaths:
+    base = pathlib.Path(base_dir)
+    base.mkdir(parents=True, exist_ok=True)
+    if run_id is None:
+        run_id = time.strftime("%Y%m%d-%H%M%S")
+    run_dir = base / run_id
+    run_dir.mkdir(parents=True, exist_ok=False)
+    return RunPaths(
+        run_dir=run_dir,
+        config_path=run_dir / "config.yaml",
+        metrics_path=run_dir / "metrics.json",
+    )
+
+
+def cmd_preprocess(args: argparse.Namespace) -> None:
+    # Prefer failing fast on missing ScenarioNet without importing heavy deps.
+    try:
+        import scenarionet  # noqa: F401
+    except ModuleNotFoundError as e:
+        raise e
+
+    from scenestreamer.dataset.dataset import SceneStreamerDataset
+
+    cfg = load_yaml_config(args.config)
+    apply_overrides(cfg, args.set or [])
+
+    # Paths: prefer CLI args, but allow config overrides.
+    if args.train_dir:
+        cfg.DATA.TRAINING_DATA_DIR = args.train_dir
+    if args.test_dir:
+        cfg.DATA.TEST_DATA_DIR = args.test_dir
+
+    cfg.DATA.USE_CACHE = True
+
+    run = make_run_dir(args.artifacts_dir, args.run_id)
+    with open(run.config_path, "w") as f:
+        yaml.safe_dump(_to_plain(cfg), f, sort_keys=False)
+
+    mode = args.split
+    ds = SceneStreamerDataset(cfg, mode)
+
+    # Iterate to materialize cache files.
+    for i in range(len(ds)):
+        _ = ds[i]
+        if args.limit is not None and (i + 1) >= args.limit:
+            break
+
+    metrics = {
+        "status": "ok",
+        "mode": mode,
+        "train_dir": getattr(cfg.DATA, "TRAINING_DATA_DIR", None),
+        "test_dir": getattr(cfg.DATA, "TEST_DATA_DIR", None),
+        "limit": args.limit,
+    }
+    with open(run.metrics_path, "w") as f:
+        json.dump(metrics, f, indent=2)
+
+    print(str(run.run_dir))
+
+
+def _load_model_from_args(args: argparse.Namespace):
+    import torch
+
+    from scenestreamer.utils import utils
+
+    device = torch.device(args.device)
+    if args.hf_repo:
+        return utils.get_model(huggingface_repo=args.hf_repo, huggingface_file=args.hf_file, device=device)
+    if args.ckpt:
+        return utils.get_model(checkpoint_path=args.ckpt, device=device)
+    raise ValueError("Must provide either --hf-repo/--hf-file or --ckpt")
+
+
+def cmd_table1(args: argparse.Namespace) -> None:
+    from scenestreamer.paper.table1_mmd import run_table1_mmd
+
+    pl_model = _load_model_from_args(args)
+    run_dir = run_table1_mmd(
+        pl_model=pl_model,
+        dataset_dir=args.dataset_dir,
+        split=args.split,
+        limit=args.limit,
+        artifacts_dir=args.artifacts_dir,
+        run_id=args.run_id,
+        seed=args.seed,
+    )
+    print(str(run_dir))
+
+
+def cmd_table2(args: argparse.Namespace) -> None:
+    from scenestreamer.paper.table2_motion import run_table2_motion
+
+    pl_model = _load_model_from_args(args)
+    run_dir = run_table2_motion(
+        pl_model=pl_model,
+        dataset_dir=args.dataset_dir,
+        split=args.split,
+        mode=args.mode,
+        num_modes=args.num_modes,
+        limit=args.limit,
+        artifacts_dir=args.artifacts_dir,
+        run_id=args.run_id,
+        seed=args.seed,
+    )
+    print(str(run_dir))
+
+
+def cmd_densify_demo(args: argparse.Namespace) -> None:
+    from scenestreamer.paper.densify_demo import run_densify_demo
+
+    pl_model = _load_model_from_args(args)
+    run_dir = run_densify_demo(
+        pl_model=pl_model,
+        dataset_dir=args.dataset_dir,
+        split=args.split,
+        scenario_index=args.scenario_index,
+        max_agents=args.max_agents,
+        force_no_end=args.force_no_end,
+        artifacts_dir=args.artifacts_dir,
+        run_id=args.run_id,
+        seed=args.seed,
+    )
+    print(str(run_dir))
+
+def _run_module_as_main(module: str, argv: list[str]) -> None:
+    old_argv = sys.argv[:]
+    try:
+        sys.argv = [module] + argv
+        runpy.run_module(module, run_name="__main__")
+    finally:
+        sys.argv = old_argv
+
+
+def cmd_table3_train(args: argparse.Namespace) -> None:
+    _run_module_as_main("scenestreamer.rl_train.train.train_td3", args.args)
+
+
+def cmd_table3_eval(args: argparse.Namespace) -> None:
+    _run_module_as_main("scenestreamer.rl_train.train.eval_policy", args.args)
+
+
+def build_parser() -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser(prog="scenestreamer", description="SceneStreamer paper reproduction CLI")
+    sub = parser.add_subparsers(dest="cmd", required=True)
+
+    def add_run_args(p: argparse.ArgumentParser) -> None:
+        p.add_argument("--artifacts-dir", default="artifacts", help="Directory to write run artifacts")
+        p.add_argument("--run-id", default=None, help="Run ID (default: timestamp)")
+        p.add_argument("--seed", type=int, default=0)
+
+    def add_model_args(p: argparse.ArgumentParser) -> None:
+        p.add_argument("--device", default="cuda", help="torch device string, e.g. cuda or cpu")
+        p.add_argument("--ckpt", default=None, help="Path to a .ckpt checkpoint")
+        p.add_argument("--hf-repo", default=None, help="HuggingFace repo id, e.g. user/repo")
+        p.add_argument("--hf-file", default=None, help="HuggingFace filename, e.g. model.ckpt")
+
+    # preprocess
+    p = sub.add_parser("preprocess", help="Preprocess ScenarioNet SD dataset and build cache")
+    add_run_args(p)
+    p.add_argument("--config", default="cfgs/motion_default.yaml")
+    p.add_argument("--set", action="append", default=[], help="Override config KEY=VALUE (repeatable)")
+    p.add_argument("--train-dir", default=None)
+    p.add_argument("--test-dir", default=None)
+    p.add_argument("--split", choices=["training", "test"], default="training")
+    p.add_argument("--limit", type=int, default=None)
+    p.set_defaults(func=cmd_preprocess)
+
+    # table1
+    p = sub.add_parser("table1", help="Table 1: initial state MMD (strict + relaxed)")
+    add_run_args(p)
+    add_model_args(p)
+    p.add_argument("--dataset-dir", required=True)
+    p.add_argument("--split", choices=["training", "test"], default="test")
+    p.add_argument("--limit", type=int, default=None)
+    p.set_defaults(func=cmd_table1)
+
+    # table2
+    p = sub.add_parser("table2", help="Table 2: motion prediction (ADE/FDE + ADD/FDD)")
+    add_run_args(p)
+    add_model_args(p)
+    p.add_argument("--dataset-dir", required=True)
+    p.add_argument("--split", choices=["training", "test"], default="test")
+    p.add_argument("--mode", choices=["motion", "full"], default="motion")
+    p.add_argument("--num-modes", type=int, default=6)
+    p.add_argument("--limit", type=int, default=None)
+    p.set_defaults(func=cmd_table2)
+
+    # demo
+    p = sub.add_parser("densify-demo", help="Qualitative densification demo (generate to max agents)")
+    add_run_args(p)
+    add_model_args(p)
+    p.add_argument("--dataset-dir", required=True)
+    p.add_argument("--split", choices=["training", "test"], default="test")
+    p.add_argument("--scenario-index", type=int, default=0)
+    p.add_argument("--max-agents", type=int, default=128)
+    p.add_argument("--force-no-end", action="store_true", help="Disable end token so it keeps generating agents")
+    p.set_defaults(func=cmd_densify_demo)
+
+    p = sub.add_parser("table3-train", help="Table 3: RL training (pass-through to train_td3.py)")
+    p.add_argument("args", nargs=argparse.REMAINDER, help="Arguments forwarded to train_td3.py")
+    p.set_defaults(func=cmd_table3_train)
+
+    p = sub.add_parser("table3-eval", help="Table 3: RL evaluation (pass-through to eval_policy.py)")
+    p.add_argument("args", nargs=argparse.REMAINDER, help="Arguments forwarded to eval_policy.py")
+    p.set_defaults(func=cmd_table3_eval)
+
+    return parser
+
+
+def main(argv: list[str] | None = None) -> None:
+    parser = build_parser()
+    args = parser.parse_args(argv)
+    try:
+        args.func(args)
+    except ModuleNotFoundError as e:
+        # Most common in a fresh environment: scenarionet / waymo-open-dataset missing.
+        msg = str(e)
+        if "scenarionet" in msg:
+            raise SystemExit(
+                "Missing dependency 'scenarionet'. Install it via:\n"
+                "  pip install git+https://github.com/metadriverse/scenarionet.git\n"
+            ) from e
+        raise
+
+
+if __name__ == "__main__":
+    main()

scenestreamer/clustering.sh

@@ -0,0 +1,7 @@
+nohup python clustering.py --data 3 > clustering_obj3_all_nomin.log 2>&1 &
+nohup python clustering.py --data 2 > clustering_obj2_all_nomin.log 2>&1 &
+nohup python clustering.py --data 1 > clustering_obj1_all_nomin.log 2>&1 &
+
+#nohup python clustering.py --data 3 --min_scale 0.5 > clustering_obj3_all.log 2>&1 &
+#nohup python clustering.py --data 2 --min_scale 0.5 > clustering_obj2_all.log 2>&1 &
+#nohup python clustering.py --data 1 --min_scale 0.5 > clustering_obj1_all.log 2>&1 &

scenestreamer/dataset/constants.py

@@ -0,0 +1,44 @@
+"""
+Define a lot of constants. It should be totally removed as most of them should be defined by MetaDrive / ScenarioNet.
+"""
+from metadrive.scenario.scenario_description import MetaDriveType
+
+# NUM_TYPES = 3
+NUM_TYPES = 5
+
+MAP_FEATURE_STATE_DIM = 27
+TRAFFIC_LIGHT_STATE_DIM = 7
+
+AGENT_STATE_DIM = 16
+
+# ACTOR_PREDICT_DIM = 6 + 2 + 4 + 5  # 3 for position, 1 for heading, 2 for velocity, 5 for types
+TRAFFIC_LIGHT_PREDICT_DIM = 9  # 9 original possible state
+
+# TODO(pzh): Do we have to do the normalization? Shouldn't the layer norm solve this?
+# POSITION_XY_RANGE = 100.
+# LOCAL_POSITION_XY_RANGE = 5.
+# HEADING_RANGE = np.pi
+# VELOCITY_XY_RANGE = 10.
+# SIZE_RANGE = 5.
+# MAP_VECTOR_XY_RANGE = 50.
+
+# TODO(pzh): Consider remove this.
+object_type_to_int = {
+    MetaDriveType.UNSET: 0,
+    MetaDriveType.VEHICLE: 1,
+    MetaDriveType.PEDESTRIAN: 2,
+    MetaDriveType.CYCLIST: 3,
+    MetaDriveType.OTHER: 4
+}
+
+# TODO(pzh): Consider remove this.
+object_int_to_type = {
+    -1: MetaDriveType.UNSET,
+    0: MetaDriveType.UNSET,
+    1: MetaDriveType.VEHICLE,
+    2: MetaDriveType.PEDESTRIAN,
+    3: MetaDriveType.CYCLIST,
+    4: MetaDriveType.OTHER
+}
+
+HEADING_PLACEHOLDER = -100  # For the object that has no heading, set this.

scenestreamer/dataset/datamodule.py

@@ -0,0 +1,49 @@
+"""
+This is a wrapper to wrap our dataset as a lightning datamodule.
+"""
+import lightning.pytorch as pl
+from torch.utils.data import DataLoader
+
+from scenestreamer.dataset import dataset
+
+
+class SceneStreamerDataModule(pl.LightningDataModule):
+    def __init__(
+        self, config, train_batch_size, train_num_workers, train_prefetch_factor, val_batch_size, val_num_workers,
+        val_prefetch_factor
+    ):
+        super().__init__()
+        self.config = config
+        self.train_batch_size = train_batch_size
+        self.train_num_workers = train_num_workers
+        self.train_prefetch_factor = train_prefetch_factor
+        self.val_batch_size = val_batch_size
+        self.val_num_workers = val_num_workers
+        self.val_prefetch_factor = val_prefetch_factor
+
+    def setup(self, stage: str):
+        self.train_dataset = dataset.SceneStreamerDataset(config=self.config, mode="training")
+        self.val_dataset = dataset.SceneStreamerDataset(config=self.config, mode="test")
+
+    def train_dataloader(self):
+        return DataLoader(
+            self.train_dataset,
+            batch_size=self.train_batch_size,
+            pin_memory=True,
+            num_workers=self.train_num_workers,
+            shuffle=True,
+            persistent_workers=True if self.train_num_workers > 0 else False,
+            collate_fn=self.train_dataset.collate_batch,
+            prefetch_factor=self.train_prefetch_factor if self.train_num_workers > 0 else None,
+        )
+
+    def val_dataloader(self):
+        return DataLoader(
+            self.val_dataset,
+            batch_size=self.val_batch_size,
+            pin_memory=True,
+            num_workers=self.val_num_workers,
+            shuffle=False,
+            collate_fn=self.val_dataset.collate_batch,
+            prefetch_factor=self.val_prefetch_factor if self.val_num_workers > 0 else None,
+        )

scenestreamer/dataset/dataset.py

@@ -0,0 +1,630 @@
+"""
+Create a pytorch dataset class for loading scenario files and padding data entries.
+"""
+import copy
+import json
+import os
+import pathlib
+import pickle
+
+try:
+    import hydra
+except ModuleNotFoundError:  # optional for core library usage
+    hydra = None
+import numpy as np
+from scenarionet import read_dataset_summary, read_scenario
+from torch.utils.data import Dataset
+
+from scenestreamer.dataset.preprocessor import preprocess_scenario_description
+from scenestreamer.utils import global_config
+from scenestreamer.utils import utils
+
+# import lmdb
+
+REPO_ROOT = pathlib.Path(__file__).resolve().parent.parent.parent
+QA_DATASET_MAPPING = {}
+ADV_INFO_DICT = {}
+
+
+class NoMapFeatureError(Exception):
+    pass
+
+
+class LMDBDatasetReader:
+    def __init__(self, base_path):
+        self.base_path = base_path
+        # Load the lookup table that maps sample keys to LMDB file names
+        # Search recursively all subfolder to find lookup.json
+        self.lookup = {}
+        for root, dirs, files in os.walk(self.base_path):
+            if "lookup.json" in files:
+                lookup_path = os.path.join(root, "lookup.json")
+                with open(lookup_path, "r") as f:
+                    lookup = json.load(f)
+                    self.lookup.update(lookup)
+        self.lmdb_cache = {}  # Cache for open LMDB environments
+
+
+#     def _get_lmdb_env(self, lmdb_name):
+#         """Fetches or opens an LMDB environment for reading."""
+#         if lmdb_name not in self.lmdb_cache:
+#             self.lmdb_cache[lmdb_name] = lmdb.open(lmdb_name, readonly=True)
+#         return self.lmdb_cache[lmdb_name]
+
+#     def load_sample(self, key):
+#         """Loads a preprocessed sample by key."""
+#         lmdb_name = self.lookup.get(key)
+#         if lmdb_name is None:
+#             raise KeyError(f"Sample {key} not found in lookup.")
+#         env = self._get_lmdb_env(lmdb_name)
+#         with env.begin() as txn:
+#             npz_bytes = txn.get(key.encode('ascii'))
+#             if npz_bytes:
+#                 with io.BytesIO(npz_bytes) as buffer:
+#                     data = np.load(buffer, allow_pickle=True)
+#                     return {name: data[name] for name in data.files}  # Return data as a dictionary
+#             return None
+
+#     def close(self):
+#         """Closes all open LMDB environments."""
+#         for env in self.lmdb_cache.values():
+#             env.close()
+
+
+def process_QA_text_label(QA_dict):
+    # TODO: do we need to form label for each individual agent? Rightnow it is just a single label
+    labels = {}
+
+    env_a = QA_dict['env_a']
+    labels['env'] = ' '.join(env_a)
+
+    ego_a = QA_dict['ego_a']
+    labels['ego'] = ' '.join(ego_a)
+
+    int_a = QA_dict['int_a']
+    labels['int'] = ' '.join(int_a)
+
+    return labels
+
+
+def get_file_paths(directory):
+    file_paths = []
+    # Traverse the directory
+    for root, dirs, files in os.walk(directory):
+        for file in files:
+            # Get the full path and add it to the list
+            full_path = os.path.join(root, file)
+            file_paths.append(full_path)
+    return
+
+
+def load_json_to_dict(file_path):
+    """
+    Load a JSON file into a Python dictionary.
+
+    :param file_path: Path to the JSON file
+    :return: Dictionary containing the JSON data
+    """
+    try:
+        with open(file_path, 'r') as file:
+            data = json.load(file)
+        return data
+    except FileNotFoundError:
+        print(f"Error: The file at {file_path} was not found.")
+    except json.JSONDecodeError:
+        print(f"Error: The file at {file_path} is not a valid JSON file.")
+    except Exception as e:
+        print(f"An unexpected error occurred: {e}")
+
+    return None
+
+
+class SceneStreamerDataset(Dataset):
+    """
+    SceneStreamer dataset class. Returns data_dict for each scenario.
+    Init args:
+        mode: "training" or "test".
+        config: 
+            - model: Details about the model architecture.
+            - data: Data directories, sample intervals, number of agents, etc.
+            - evaluation: predict_all_agents, delete_eval_result (TODO: Add ScenarioDescription passthrough as a flag in the config.)
+            - optimization: Training hyperparameters.
+            - preprocessing: Max number of agents, map features, traffic lights, padding, etc.
+            - root_dir: Self-explanatory.
+            - sampling: Inference sampling parameters.
+            - tokenization: The part of the config passed to the tokenizer.
+    """
+    def __init__(self, config, mode):
+        super().__init__()
+        self.mode = mode
+        self.config = config
+        dataset_cfg = self.config.DATA
+
+        self.max_map_features = config.PREPROCESSING.MAX_MAP_FEATURES
+        self.max_vectors_per_map_feature = config.PREPROCESSING.MAX_VECTORS
+        self.max_agents = config.PREPROCESSING.MAX_AGENTS
+        self.max_traffic_lights = config.PREPROCESSING.MAX_TRAFFIC_LIGHTS
+        self.padding_to_max = config.PREPROCESSING.PADDING_TO_MAX
+
+        # We are expecting the data_dir to be either an absolute path or a relative path w.r.t. the repo root.
+        if mode == "training":
+            self.data_dir = global_config.ROOT_DIR / dataset_cfg.TRAINING_DATA_DIR
+        elif mode == "test":
+            self.data_dir = global_config.ROOT_DIR / dataset_cfg.TEST_DATA_DIR
+        else:
+            raise ValueError(f"Unknown mode {mode}.")
+
+        # summary_dict: A dictionary of .pkl filenames to ingest. Filenames (keys) are mapped to metadata objects.
+        # summary_list: Keys of summary_dict, in order of ingestion.
+        # mapping: A dict mapping scenario IDs to the folder that hosts their files.
+        summary_dict, summary_list, mapping = read_dataset_summary(self.data_dir)
+
+        # We might want to use a subset of scenarios.
+        if self.mode == "training":
+            interval = dataset_cfg.SAMPLE_INTERVAL_TRAINING
+        elif self.mode == "test":
+            interval = dataset_cfg.SAMPLE_INTERVAL_TEST
+        else:
+            raise ValueError(f"Unknown mode {self.mode}.")
+
+        if "SD_PASSTHROUGH" in config.DATA:
+            self.return_scenario_description = config.DATA["SD_PASSTHROUGH"]
+        else:  # Default to False.
+            self.return_scenario_description = False
+
+        summary_list = summary_list[::interval]
+        # self.data_summary_dict = {k: summary_dict[k] for k in summary_list}
+        self.data_mapping = {k: mapping[k] for k in summary_list}
+        self.length = len(summary_list)
+        self.use_cache_logged = False
+
+        if self.config.BACKWARD_PREDICTION and self.mode == "training":
+            self.real_length = self.length
+            self.length = self.length * 2
+
+        # Convert each string to sequence of codepoints (integer),
+        # and then pack them into a numpy array.
+        # NOTE(pzh): I forgot why I wrote this. Seems like some issues in multiprocessing.
+
+        # seqs: A list of np.arrays, each representing the ascii values of a string.
+        seqs = [utils.string_to_sequence(s) for s in summary_list]
+
+        # strings_v: ascii values of all strings, concatenated.
+        # strings_o: offsets of each string in strings_v.
+        if len(seqs) == 0:
+            raise ValueError("No scenarios found in the dataset: {}".format(self.data_dir))
+        self.strings_v, self.strings_o = utils.pack_sequences(seqs)
+
+        # if self.config.DATA.USE_LMDB and self.mode == "training":
+        #     cache_folder = pathlib.Path(self.data_dir) / "cache"
+        #     assert cache_folder.is_dir()
+        #     self.reader = LMDBDatasetReader(cache_folder)  # LMDB Reader to load samples
+
+        from scenestreamer.tokenization import get_tokenizer
+        self.tokenizer = get_tokenizer(config=self.config)
+
+    def __len__(self):
+        return self.length
+
+    def __getitem__(self, index):
+        # Unpack the stored codepoints at the correct index into a filename string.
+
+        use_backward_prediction = False
+        if self.config.BACKWARD_PREDICTION and self.mode == "training":
+            if index >= self.real_length:
+                index = index - self.real_length
+                use_backward_prediction = True
+
+        seq = utils.unpack_sequence(self.strings_v, self.strings_o, index)
+        string = utils.sequence_to_string(seq)
+        file_name = string
+
+        try:
+            data_dict = self.create_scene_level_data(file_name, index, use_backward_prediction)
+        except NoMapFeatureError:
+            # This is workaround for Waymo test set where some scenarios do not have map features.
+            return self.__getitem__(index + 1)
+
+        # If self.return_scenario_description is true, data_dict has an extra key [raw_scenario_description] that contains the ScenarioDescription object.
+        return data_dict
+
+    def create_scene_level_data(self, file_name, index, use_backward_prediction=False):
+        """
+        Reads a scenario file and preprocesses it.
+        """
+        assert not self.config.DATA.USE_LMDB, "LMDB is not supported."
+        try:
+            # scenario: A ScenarioDescription instance.
+            cache = None
+            scenario = None
+            cache_path = None
+
+            if self.config.DATA.USE_CACHE:
+                cache_folder = pathlib.Path(self.data_dir) / "cache"
+                if cache_folder.is_dir() is False:
+                    cache_folder.mkdir(exist_ok=True)
+
+                cache_path = pathlib.Path(self.data_dir) / "cache" / file_name
+                if cache_path.is_file():
+
+                    try:
+                        with open(cache_path, "rb") as f:
+                            cache = pickle.load(f)
+
+                        if self.use_cache_logged is False:
+                            print("=====================================")
+                            print("=====================================")
+                            print("\t*** WARNING ***")
+                            print("\tYou are using cache files!!!")
+                            print("\tIn folder: ", cache_folder)
+                            print("\tThere are ", len(list(cache_folder.glob("*"))), " cache files!!!")
+                            print("=====================================")
+                            print("=====================================")
+
+                            self.use_cache_logged = True
+
+                        return cache
+                    except EOFError as e:
+                        print(f"Error in reading cache file: {cache_path=}")
+
+                    scenario = read_scenario(
+                        dataset_path=self.data_dir, mapping=self.data_mapping, scenario_file_name=file_name
+                    )
+
+                else:
+                    scenario = read_scenario(
+                        dataset_path=self.data_dir, mapping=self.data_mapping, scenario_file_name=file_name
+                    )
+                    # print("Cannot find cache file: ", cache_path, "Creating one.")
+
+            else:
+                # if self.config.DATA.USE_LMDB and self.mode == "training":
+                #     cache = self.reader.load_sample(file_name)
+                # else:
+                scenario = read_scenario(
+                    dataset_path=self.data_dir, mapping=self.data_mapping, scenario_file_name=file_name
+                )
+
+        except EOFError as e:
+            print(f"{self.data_dir=}, {self.data_mapping=}, {file_name=}")
+            raise e
+        assert self.mode in ["training", "test"], self.mode
+        ret = {}
+
+        if len(scenario["map_features"]) == 0:
+            raise NoMapFeatureError
+
+        if self.return_scenario_description:
+            ret["raw_scenario_description"] = copy.deepcopy(scenario)
+
+        # TODO: Remove error handling after debugging.
+        try:
+            preprocessed_scenario_description = preprocess_scenario_description(
+                scenario=scenario,
+                # cache=cache,
+                config=copy.deepcopy(self.config),
+                in_evaluation=self.mode != "training",
+                keep_all_data=self.config.PREPROCESSING.get("keep_all_data", False),
+                backward_prediction=use_backward_prediction,
+                tokenizer=self.tokenizer,
+                # cache_path=cache_path,
+            )
+            preprocessed_scenario_description["file_name"] = file_name
+        except Exception as e:
+            print(f"Error in preprocessing {file_name=}, {index=}, {scenario['id']=}")
+            # Ensure that the exception is not swallowed by adding this.
+            raise RuntimeError(
+                f"{file_name=}, {index=}, {scenario['id']=}. Error in create_scene_level_data: {e}"
+            ) from e
+
+        ret.update(preprocessed_scenario_description)
+        ret.update({"metadata/scenario_id": scenario['id']})
+
+        if cache_path is not None:
+            with open(cache_path, "wb") as f:
+                pickle.dump(ret, f)
+                # print("Writing cache file: ", cache_path)
+
+        return ret
+
+    def collate_batch(self, batch_list):
+        """
+        Output format:
+
+        agent_feature:              [B, T, #agents, D]
+        agent_feature_position:     [B, T, #agents, 3]
+        map_feature:                [B, T, #mapfeat, #points, D]
+        map_feature_valid_mask:     [B, T, #mapfeat, #points]
+        map_feature_position:       [B, T, #mapfeat, 3]
+        """
+        data_dict_sample = batch_list[0]
+
+        num_map_feat, num_points, _ = data_dict_sample["encoder/map_feature"].shape
+
+        data_dict = {}
+        object_keys = [
+            "raw_scenario_description",
+            "encoder/track_name",
+            "decoder/track_name",
+            "eval/track_name",
+            # "scenario_id",
+            # "in_evaluation"
+        ]  # Keys exempt from padding and tensor conversion.
+
+        for k in set(data_dict_sample.keys()):
+            if k not in object_keys:
+                if not isinstance(data_dict_sample[k], np.ndarray):
+                    assert isinstance(data_dict_sample[k], (int, float, bool, str)), (k, type(data_dict_sample[k]))
+                    if isinstance(data_dict_sample[k], str):
+                        data_dict[k] = np.array([b[k] for b in batch_list])
+                    else:
+                        data_dict[k] = utils.numpy_to_torch(np.array([b[k] for b in batch_list]))
+                    continue
+                # else:
+                #     if batch_list[0][k].dtype == np.object:
+                #         data_dict[k] = [b[k] for b in batch_list]
+                #         continue
+
+                val_list = [utils.numpy_to_torch(b[k]) for b in batch_list]
+
+            # Map features that have vectors' information
+            if k in [
+                    "encoder/map_feature",
+                    "vis/map_feature",
+                    "raw/map_feature",
+                    "encoder/map_feature_valid_mask",
+            ]:
+                data_dict[k] = utils.padding_1st_and_2nd_dim(
+                    val_list,
+                    max_1st_dim=self.max_map_features if self.padding_to_max else None,
+                    max_2nd_dim=self.max_vectors_per_map_feature if self.padding_to_max else None
+                )
+
+            # Map features that have aggregated info from vectors
+            elif k in [
+                    "encoder/map_heading",
+                    "encoder/map_position",
+                    "encoder/map_valid_mask",
+            ]:
+                data_dict[k] = utils.padding_1st_dim(
+                    val_list, max_1st_dim=self.max_map_features if self.padding_to_max else None
+                )
+
+            # Traffic light features that have temporal dim
+            elif k in [
+                    "encoder/traffic_light_feature",
+                    "encoder/traffic_light_state",
+                    "encoder/traffic_light_valid_mask",
+            ]:
+
+                if self.config.PREPROCESSING.REMOVE_TRAFFIC_LIGHT_STATE:
+                    data_dict[k] = utils.padding_1st_dim(
+                        val_list, max_1st_dim=self.max_traffic_lights if self.padding_to_max else None
+                    )
+                else:
+                    data_dict[k] = utils.padding_1st_and_2nd_dim(
+                        val_list, max_2nd_dim=self.max_traffic_lights if self.padding_to_max else None
+                    )
+
+            # Traffic light features that do not have temporal dim
+            elif k in [
+                    "encoder/traffic_light_position",
+                    "encoder/traffic_light_heading",
+                    "encoder/traffic_light_map_id",
+            ]:
+                data_dict[k] = utils.padding_1st_dim(
+                    val_list, max_1st_dim=self.max_traffic_lights if self.padding_to_max else None
+                )
+
+            # Agent features
+            elif k in [
+                    "encoder/agent_feature",
+                    "encoder/agent_position",
+                    "encoder/agent_valid_mask",
+                    "encoder/agent_heading",
+                    "encoder/agent_velocity",
+                    "decoder/modeled_agent_position",
+                    "decoder/modeled_agent_heading",
+                    "decoder/modeled_agent_velocity",
+                    "decoder/modeled_agent_delta",
+            ]:
+                data_dict[k] = utils.padding_1st_and_2nd_dim(
+                    val_list, max_2nd_dim=self.max_agents if self.padding_to_max else None
+                )
+
+            # Other data that does not pass the model or does not need regular shapes
+            elif k in [
+                    # "encoder/modeled_agent_id",
+                    # "action_label/labeled_agent_id",
+                    "metadata/map_center",  # "decoder/input_step",
+                    # "decoder/input_intra_step",
+                    "encoder/current_agent_heading",
+                    "decoder/current_agent_heading",
+                    "encoder/current_agent_shape",
+                    "decoder/current_agent_shape",
+                    "eval/current_agent_heading",
+                    "encoder/current_agent_valid_mask",
+                    "decoder/current_agent_valid_mask",
+                    "eval/current_agent_valid_mask",
+                    # "decoder/current_agent_valid_mask",  #
+                    # "decoder/modeled_agent_indices",
+                    # For gen model:
+                    # "decoder/input_token_valid_mask",
+                    # "decoder/should_predict",
+                    # "decoder/is_gt",
+                    # "eval/should_predict_motion",
+            ]:
+                data_dict[k] = utils.padding_1st_dim(val_list)
+
+            elif k in [
+                    "decoder/input_action_valid_mask",
+                    "encoder/current_agent_position",
+                    "decoder/current_agent_position",
+                    "encoder/current_agent_velocity",
+                    "decoder/current_agent_velocity",
+                    "decoder/target_action_valid_mask",
+                    #"decoder/future_agent_position",
+                    #"decoder/future_agent_heading",
+                    #"decoder/future_agent_valid_mask",
+                    #"decoder/future_agent_velocity",
+                    #"encoder/future_agent_position",
+                    #"encoder/future_agent_heading",
+                    #"encoder/future_agent_valid_mask",
+                    #"encoder/future_agent_velocity",
+                    "decoder/agent_position",
+                    "decoder/agent_heading",
+                    "decoder/agent_velocity",
+                    "decoder/agent_valid_mask",
+                    "eval/agent_velocity",
+                    "eval/agent_heading",
+                    "eval/agent_position",
+                    "eval/agent_valid_mask",
+                    "encoder/agent_shape",
+                    "decoder/agent_shape",
+                    "eval/agent_shape",  # "decoder/target_valid_mask",
+                    "decoder/input_agent_motion",
+                    "decoder/target_agent_motion",
+                    "decoder/dest_map_index_valid_mask",
+            ]:
+                data_dict[k] = utils.padding_1st_and_2nd_dim(val_list)
+
+            elif k in [
+                    "encoder/agent_type",
+                    "decoder/agent_type",
+                    "encoder/modeled_agent_type",
+                    "eval/agent_type",  # "eval/raw_agent_name",
+                    "encoder/object_of_interest_name",
+                    "decoder/object_of_interest_name",
+                    "metadata/sdc_name",  # "eval/modeled_agent_id",
+                    "encoder/object_of_interest_id",
+                    "decoder/object_of_interest_id",
+                    "encoder/modeled_agent_id",  # "decoder/modeled_agent_id",
+                    "encoder/agent_id",
+                    "decoder/agent_id",
+                    "decoder/labeled_agent_id",
+                    "decoder/label_turning",
+                    "decoder/label_acceleration",
+                    "decoder/label_safety",
+                    # For gen model:
+                    #                   "decoder/input_token_id",
+                    #                   "decoder/causal_mask_offset",
+            ]:
+                data_dict[k] = utils.padding_1st_dim(val_list, fill=-1)
+
+            elif k in [
+                    "decoder/dest_map_index",
+                    "decoder/dest_map_index_gt",
+            ]:
+                data_dict[k] = utils.padding_1st_and_2nd_dim(val_list, fill=-1)
+
+            elif k in [
+                    "decoder/input_action",
+                    "decoder/target_action",
+                    "decoder/input_action_for_trafficgen",
+
+                "decoder/current_agent_shape_for_trafficgen",
+                "decoder/modeled_agent_heading_for_trafficgen",
+                "decoder/modeled_agent_position_for_trafficgen",
+                "decoder/modeled_agent_velocity_for_trafficgen",
+                "decoder/input_action_valid_mask_for_trafficgen",
+                "decoder/modeled_agent_delta_for_trafficgen",
+                "decoder/input_action_feature_for_trafficgen",
+                "decoder/target_offset_for_trafficgen",
+                "decoder/input_offset_for_trafficgen",
+                "decoder/agent_id_for_trafficgen",
+                "decoder/trafficgen_position",
+                "decoder/trafficgen_heading",
+                "decoder/agent_type_for_trafficgen",
+            ]:
+                data_dict[k] = utils.padding_all_dims(val_list, fill=-1)
+
+            elif k in object_keys:
+                # Passthrough: Have the data_dict[object] contain a list of objects.
+                data_dict[k] = [b[k] for b in batch_list]
+
+            elif k in [
+                    "encoder/sdc_index",
+            ]:
+                pass
+
+            else:
+                raise ValueError("Unknown key: {}".format(k))
+
+        return data_dict
+
+
+if hydra is not None:
+    @hydra.main(version_base=None, config_path=str(REPO_ROOT / "cfgs"), config_name="1009_safety_action_debug.yaml")
+    def debug(config):
+        test_dataset = SceneStreamerDataset(config, "training")
+        ddd = iter(test_dataset)
+        count = 0
+        buggy_count = 0
+        while True:
+            if count == 3:
+                return
+            try:
+                data = next(ddd)
+                count += 1
+
+                assert data["decoder/label_safety"][data["decoder/labeled_agent_id"]].sum() > 1
+
+            except StopIteration:
+                break
+
+            except AssertionError:
+                print("ni collision")
+                buggy_count += 1
+                print("scenario_id", data["scenario_id"])
+                print("data['decoder/label_safety']", data["decoder/label_safety"])
+                print("data['decoder/labeled_agent_id']", data["decoder/labeled_agent_id"])
+                print("track_name", data["decoder/track_name"][data["decoder/labeled_agent_id"]])
+
+        print("buggy_count:", buggy_count)
+        print("count", count)
+        print("End")
+
+    @hydra.main(version_base=None, config_path=str(REPO_ROOT / "cfgs"), config_name="motion_default.yaml")
+    def read_traffic_light_state(config):
+        test_dataset = SceneStreamerDataset(config, "training")
+
+        total_tl = 0
+        total_green = 0
+        total_yellow = 0
+        total_red = 0
+        total_unknown = 0
+        total_mix = 0
+        import tqdm
+
+        for data in tqdm.tqdm(test_dataset):
+            tl = data["encoder/traffic_light_feature"]
+            mask = data["encoder/traffic_light_valid_mask"]
+
+            for i in range(tl.shape[1]):
+                if mask[:, i].any():
+                    is_green = tl[:, i, 3].astype(bool).any()
+                    is_yellow = tl[:, i, 4].astype(bool).any()
+                    is_red = tl[:, i, 5].astype(bool).any()
+                    is_unknown = tl[:, i, 6].astype(bool).any()
+
+                    total_tl += 1
+                    total_green += is_green
+                    total_yellow += is_yellow
+                    total_red += is_red
+                    total_unknown += is_unknown
+                    total_mix += (is_green and is_yellow) or (is_green and is_red) or (is_yellow and is_red)
+
+        print("total_tl:", total_tl)
+        print("total_green: {}\t{:.4f}".format(total_green, total_green / total_tl))
+        print("total_yellow: {}\t{:.4f}".format(total_yellow, total_yellow / total_tl))
+        print("total_red: {}\t{:.4f}".format(total_red, total_red / total_tl))
+        print("total_unknown: {}\t{:.4f}".format(total_unknown, total_unknown / total_tl))
+        print("total_mix: {}\t{:.4f}".format(total_mix, total_mix / total_tl))
+else:
+    debug = None
+    read_traffic_light_state = None
+
+
+if __name__ == '__main__':
+    # debug()
+    read_traffic_light_state()

scenestreamer/dataset/make_lmdb.py

@@ -0,0 +1,233 @@
+"""
+Only the TRAINING_DATA_DIR will be used in the code below.
+Usage:
+
+python -m scenestreamer.dataset.make_lmdb \
+--config-name="1024_gpt" DATA.TEST_DATA_DIR='data/20scenarios' \
+DATA.TRAINING_DATA_DIR="/data_zhenghao/datasets/scenarionet/CAT_waymo_hybrid/"
+
+"""
+import json
+import os
+import pathlib
+import pickle
+import multiprocessing as mp
+from functools import partial
+import tqdm
+import hydra
+import lmdb
+import omegaconf
+import tqdm
+
+from scenestreamer.dataset.dataset import SceneStreamerDataset
+
+REPO_ROOT = pathlib.Path(__file__).resolve().parent.parent.parent
+
+
+class LMDBBulkWriter:
+    def __init__(self, base_path, max_size=1e9):
+        """
+        Initializes the LMDBBulkWriter to save all data in batches, with map_size for each LMDB file.
+        Args:
+            base_path: Directory path to save LMDB files.
+            max_size: Maximum size of each LMDB file in bytes.
+        """
+        self.base_path = base_path
+        # Create the cache directory if it doesn't exist
+        os.makedirs(self.base_path, exist_ok=True)
+
+        self.max_size = int(max_size)  # Set the max LMDB file size (e.g., 1 GB)
+        self.current_db_index = 0
+        self.lookup = {}  # Lookup table to track which LMDB file stores which sample
+        self.current_db = self._open_new_lmdb(self.current_db_index)
+        self.per_shard_size = 0
+
+        self.sample_buffer = []
+
+    def _open_new_lmdb(self, db_index):
+        """Opens a new LMDB file for saving samples."""
+        db_path = f"{self.base_path}/data_{db_index}.lmdb"
+        return lmdb.open(db_path, map_size=self.max_size)
+
+    def _save_a_batch(self):
+
+        try:
+            # Commit the transaction if we have reached the commit interval
+            # if (not hasattr(self, 'txn')) or (self.txn is None):
+            #     self.txn = self.current_db.begin(write=True)  # Start a new transaction
+            #
+            #
+            # for key, data in self.sample_buffer:
+            #     self.txn.put(key.encode('ascii'), pickle.dumps(data))
+            #     self.lookup[key] = f"data_{self.current_db_index}.lmdb"
+            #
+            #     if hasattr(self, 'txn') and self.txn:
+            #         self.txn.commit()  # Commit the transaction
+            print(f"Saving {len(self.sample_buffer)} samples to data_{self.current_db_index}.lmdb")
+            with self.current_db.begin(write=True) as txn:
+                for key, data in self.sample_buffer:
+                    txn.put(key.encode('ascii'), pickle.dumps(data))
+                    self.lookup[key] = f"data_{self.current_db_index}.lmdb"
+
+            self.sample_buffer.clear()
+
+        except lmdb.MapFullError:
+
+            # If current LMDB file is full, create a new one and retry saving
+            self.current_db.close()
+            self.current_db_index += 1
+            print(f"Creating new LMDB file: data_{self.current_db_index}.lmdb (size: {self.per_shard_size})")
+            self.current_db = self._open_new_lmdb(self.current_db_index)
+            self._save_a_batch()
+            self.per_shard_size = 0
+
+    def save_sample(self, key, data):
+        """Saves a sample to the current LMDB file, switching to a new file if necessary."""
+        # Batch writes into a single transaction
+        if self.per_shard_size % 100 == 0:
+            self._save_a_batch()
+        self.sample_buffer.append((key, data))
+        self.per_shard_size += 1
+
+    def close(self):
+        self._save_a_batch()
+        """Closes the LMDB environment and saves the lookup table as a JSON file."""
+        self.current_db.close()
+        # Save the lookup table to track the LMDB file where each sample is stored
+        with open(f"{self.base_path}/lookup.json", "w") as f:
+            json.dump(self.lookup, f)
+
+
+def preprocess_and_queue_worker(worker_id, config, indices, queue):
+    """
+    This function runs in each worker to preprocess samples and send them to the write queue.
+    The writer process will handle writing to LMDB.
+    """
+    print(f"Worker {worker_id} started.")
+    dataset = SceneStreamerDataset(config, "training")
+
+    print(f"Worker {worker_id} has {len(dataset)} samples.")
+
+    # Process and queue each sample assigned to this worker
+    if worker_id == 0:
+        pbar = tqdm.tqdm(indices, desc="Worker %d" % worker_id)
+    else:
+        pbar = indices
+    print(f"Worker {worker_id} has {len(indices)} samples.")
+
+    for i in pbar:
+        sample = dataset[i]  # Access the sample using its index
+
+        # Simulate some preprocessing (replace with actual preprocessing logic)
+        file_name, processed_sample = sample["file_name"], sample
+
+        # Put the preprocessed sample into the queue to be written by the writer process
+        print(f"Worker {worker_id} processed {file_name}")
+        queue.put((file_name, processed_sample))
+
+    # Signal that this worker is done
+    # queue.put(None)  # 'None' signals that the worker is done
+
+
+def write_process(queue, base_path, max_size):
+    """
+    The write process receives samples from the queue and writes them to the LMDB environment.
+    """
+    writer = LMDBBulkWriter(base_path=base_path, max_size=max_size)
+    print("Writer process started.")
+
+    while True:
+
+        # Blocking if no data is available
+        data = queue.get()
+
+        if data == 100:
+            print("Received 100, stopping writer process.")
+            # If 'None' is received, this indicates that a worker has finished
+            break
+
+        if data is None:
+            print("Received None, stopping writer process.")
+            continue
+
+        file_name, processed_sample = data
+        print(f"Saved {file_name} to LMDB")
+        writer.save_sample(file_name, processed_sample)
+
+    # Close the writer once all workers are done
+    writer.close()
+
+
+@hydra.main(version_base=None, config_path=str(REPO_ROOT / "cfgs"), config_name="1024_gpt.yaml")
+def make_lmdb(config):
+    omegaconf.OmegaConf.set_struct(config, False)
+    omegaconf.OmegaConf.set_struct(config, True)
+
+    dataset = SceneStreamerDataset(config, "training")
+    folder = pathlib.Path(dataset.data_dir)
+    folder = folder / "cache"
+    folder.mkdir(parents=True, exist_ok=False)
+
+    # Initialize the LMDBBulkWriter
+    print("Saving data to LMDB folder:", folder.absolute())
+
+    # num_workers = mp.cpu_count()
+    num_workers = 2
+
+    dataset_size = len(dataset)
+    indices = list(range(dataset_size))
+    chunk_size = dataset_size // num_workers
+
+    # Split the indices into chunks, one for each worker
+    chunked_indices = [indices[i * chunk_size:(i + 1) * chunk_size] for i in range(num_workers)]
+    # The final chunk may have more samples if the dataset size is not divisible by the number of workers.
+    chunked_indices[0].extend(indices[num_workers * chunk_size:])
+
+    # Create a multiprocessing queue
+    queue = mp.Queue()
+
+    # Create and start the writer process
+    writer_process = mp.Process(target=write_process, args=(queue, folder, 1e10))
+
+    writer_process.start()
+
+    # Create a multiprocessing pool for parallel processing (preprocessing)
+    pool = mp.Pool(num_workers)
+
+    results = []
+    # Start each worker process, passing its chunk of indices
+    for worker_id, worker_indices in enumerate(chunked_indices):
+        print(f"Starting worker {worker_id} with {len(worker_indices)} samples.")
+        result = pool.apply_async(preprocess_and_queue_worker, args=(worker_id, config, worker_indices, queue))
+        results.append(result)
+
+    # Wait for all worker processes to complete
+    # for result in results:
+    #     result.get()  # This will block until the worker completes its task
+    # preprocess_and_queue_worker(0, config, chunked_indices[0], queue)
+    pool.close()
+    print("Waiting for workers to finish...")
+    pool.join()
+    print("All workers finished.")
+
+    # Signal the writer process to stop (send 'None' once all workers are done)
+    queue.put(100)
+    # Wait for the writer process to finish
+    writer_process.join()
+
+
+@hydra.main(version_base=None, config_path=str(REPO_ROOT / "cfgs"), config_name="motion_default.yaml")
+def debug(config):
+    omegaconf.OmegaConf.set_struct(config, False)
+    omegaconf.OmegaConf.set_struct(config, True)
+    dataset = SceneStreamerDataset(config, "training")
+    folder = pathlib.Path(dataset.data_dir)
+    folder = folder / "cache"
+    folder.mkdir(parents=True, exist_ok=True)
+    for i, sample in enumerate(tqdm.tqdm(dataset, total=len(dataset), desc="Scenarios")):
+        file_name = sample["file_name"]
+
+
+if __name__ == '__main__':
+    make_lmdb()
+    # debug()

scenestreamer/dataset/preprocess_action_label.py

@@ -0,0 +1,293 @@
+import numpy as np
+from shapely.geometry import Polygon
+
+from scenestreamer.utils import utils
+
+INVALID_VALUE = -10000
+
+
+class TurnAction:
+    STOP = 0
+    KEEP_STRAIGHT = 1
+    TURN_LEFT = 2
+    TURN_RIGHT = 3
+    U_TURN = 4
+
+    num_actions = 5
+
+
+class AccelerationAction:
+    STOP = 0
+    KEEP_SPEED = 1
+    SPEED_UP = 2
+    SLOW_DOWN = 3
+
+    num_actions = 4
+
+
+class SafetyAction:
+    SAFE = 0
+    COLLISION = 1
+    num_actions = 2
+
+
+def cal_polygon_contour(x, y, theta, width, length):
+
+    left_front_x = x + 0.5 * length * np.cos(theta) - 0.5 * width * np.sin(theta)
+    left_front_y = y + 0.5 * length * np.sin(theta) + 0.5 * width * np.cos(theta)
+    left_front = np.column_stack((left_front_x, left_front_y))
+
+    right_front_x = x + 0.5 * length * np.cos(theta) + 0.5 * width * np.sin(theta)
+    right_front_y = y + 0.5 * length * np.sin(theta) - 0.5 * width * np.cos(theta)
+    right_front = np.column_stack((right_front_x, right_front_y))
+
+    right_back_x = x - 0.5 * length * np.cos(theta) + 0.5 * width * np.sin(theta)
+    right_back_y = y - 0.5 * length * np.sin(theta) - 0.5 * width * np.cos(theta)
+    right_back = np.column_stack((right_back_x, right_back_y))
+
+    left_back_x = x - 0.5 * length * np.cos(theta) - 0.5 * width * np.sin(theta)
+    left_back_y = y - 0.5 * length * np.sin(theta) + 0.5 * width * np.cos(theta)
+    left_back = np.column_stack((left_back_x, left_back_y))
+
+    polygon_contour = np.concatenate(
+        (left_front[:, None, :], right_front[:, None, :], right_back[:, None, :], left_back[:, None, :]), axis=1
+    )
+
+    return polygon_contour
+
+
+def detect_collision(contour_list1, mask1, contour_list2, mask2):
+    collision_detected = []
+    assert len(contour_list1) == len(contour_list2)
+
+    for i in range(len(contour_list1)):
+        if mask1[i] and mask2[i]:
+            poly1 = Polygon(contour_list1[i])
+            poly2 = Polygon(contour_list2[i])
+
+            if poly1.intersects(poly2):
+                collision_detected.append(True)
+            else:
+                collision_detected.append(False)
+        else:
+            collision_detected.append(False)
+
+    return collision_detected
+
+
+def get_direction_action_from_trajectory_batch(traj, mask, dt=0.1, ooi=None):
+    U_TURN_DEG = 115
+    LEFT_TURN_DEG = 25
+    RIGHT_TURN_DEG = -25
+    STOP_SPEED = 0.06
+
+    assert traj.ndim == 3
+    traj_diff = traj[1:] - traj[:-1]
+    mask_diff = mask[1:] & mask[:-1]
+
+    displacement = np.linalg.norm(traj_diff, axis=-1)
+
+    mask_diff_stop = mask_diff & (displacement > 0.1)
+
+    pred_angles = np.arctan2(traj_diff[..., 1], traj_diff[..., 0])
+    pred_angles_diff = utils.wrap_to_pi(pred_angles[1:] - pred_angles[:-1])
+
+    # It's meaning less to compute heading for a stopped vehicle. So mask them out!
+    mask_diff_diff = mask_diff_stop[1:] & mask_diff_stop[:-1]
+    # Note that we should not wrap to pi here because the sign is important.
+    accumulated_heading_change_rad = (pred_angles_diff * mask_diff_diff).sum(axis=0)
+    accumulated_heading_change_deg = np.degrees(accumulated_heading_change_rad)
+
+    # print("accumulated_heading_change_deg: ", list(zip(ooi, accumulated_heading_change_deg)))
+
+    speed = displacement / dt
+    avg_speed = utils.masked_average_numpy(speed, mask_diff, dim=0)
+
+    actions = np.zeros(accumulated_heading_change_deg.shape, dtype=int)
+    actions.fill(TurnAction.KEEP_STRAIGHT)
+    actions[accumulated_heading_change_deg > LEFT_TURN_DEG] = TurnAction.TURN_LEFT
+    actions[accumulated_heading_change_deg < RIGHT_TURN_DEG] = TurnAction.TURN_RIGHT
+    actions[accumulated_heading_change_deg > U_TURN_DEG] = TurnAction.U_TURN
+    actions[accumulated_heading_change_deg < -U_TURN_DEG] = TurnAction.U_TURN
+    actions[avg_speed < STOP_SPEED] = TurnAction.STOP
+    return actions
+
+
+def get_acce_action_from_trajectory_batch(batch_trajs, mask, ooi=None, dt=0.1):
+
+    SPEEDUP_ACCEL = 0.3
+    SPEEDDOWN_ACCEL = -0.3
+    STOP_SPEED = 0.06
+
+    traj_diff = batch_trajs[1:] - batch_trajs[:-1]  # (T, A, 2)
+    mask_diff = mask[1:] & mask[:-1]  # (T, A)
+
+    speed = np.linalg.norm(traj_diff, axis=-1) / dt  # (T, A)
+
+    speed_change = speed[1:] - speed[:-1]
+    mask_diff_diff = mask_diff[1:] & mask_diff[:-1]
+
+    absolute_avg_speed = utils.masked_average_numpy(speed, mask_diff, dim=0)
+
+    accumulated_speed_change = (speed_change * mask_diff_diff).sum(0)
+
+    init_speed_ind = mask_diff.argmax(axis=0)
+    init_speed = np.take_along_axis(speed, init_speed_ind[None, :], axis=0)[0]
+
+    speed_change_ratio = accumulated_speed_change / np.maximum(init_speed, STOP_SPEED)
+
+    # print("speed_change_ratio: ", list(zip(ooi, speed_change_ratio)))
+
+    actions = np.zeros(speed_change_ratio.shape, dtype=int)
+
+    actions.fill(AccelerationAction.KEEP_SPEED)
+    actions[speed_change_ratio > SPEEDUP_ACCEL] = AccelerationAction.SPEED_UP
+    actions[speed_change_ratio < SPEEDDOWN_ACCEL] = AccelerationAction.SLOW_DOWN
+    actions[absolute_avg_speed <= STOP_SPEED] = AccelerationAction.STOP  # if stop
+
+    return actions
+
+
+def get_safety_action_from_sdc_adv(data_dict, adv_id, sdc_id):
+
+    contours = []
+    for agent_id in [adv_id, sdc_id]:
+        traj = data_dict["decoder/agent_position"][:91, agent_id, :]  # (91, 3)
+        length = data_dict["decoder/agent_shape"][:91, agent_id, 0]
+        width = data_dict["decoder/agent_shape"][:91, agent_id, 1]
+        theta = data_dict['decoder/agent_heading'][:91, agent_id]  # (91, ) # in pi
+        mask = data_dict['decoder/agent_valid_mask'][:91, agent_id]  # (91,)
+
+        poly = cal_polygon_contour(traj[:, 0], traj[:, 1], theta, width, length)
+        contours.append(poly)
+
+    sdc_mask = data_dict['decoder/agent_valid_mask'][:, sdc_id]  # (91,)
+    adv_mask = data_dict['decoder/agent_valid_mask'][:, adv_id]
+    adv_contour = contours[0]
+    sdc_contour = contours[1]
+
+    collision_detected = detect_collision(adv_contour, adv_mask, sdc_contour, sdc_mask)
+
+    # instead of loading a dict which saves all collision scenario, we could simply detect all agents' potential collision
+    return collision_detected
+
+
+def get_safety_action_from_trajectory_batch(data_dict, track_agent_indicies):
+
+    safety_actions = np.zeros((track_agent_indicies.shape[0], ), dtype=int)  # plus sdc
+
+    contours = []
+    for agent1_id in track_agent_indicies:
+        traj = data_dict["decoder/agent_position"][:, agent1_id, :]  # (91, 3)
+        length = data_dict["decoder/agent_shape"][:, agent1_id, 0]
+        width = data_dict["decoder/agent_shape"][:, agent1_id, 1]
+        theta = data_dict['decoder/agent_heading'][:, agent1_id]  # (91, ) # in pi
+        mask = data_dict['decoder/agent_valid_mask'][:, agent1_id]  # (91,)
+        poly = cal_polygon_contour(traj[:, 0], traj[:, 1], theta, width, length)
+        contours.append(poly)
+
+    for i in range(track_agent_indicies.shape[0] - 1):
+        for j in range(i + 1, track_agent_indicies.shape[0]):
+            mask_1 = data_dict['decoder/agent_valid_mask'][:, track_agent_indicies[i]]  # (91,)
+            mask_2 = data_dict['decoder/agent_valid_mask'][:, track_agent_indicies[j]]
+            collision_detected = detect_collision(contours[i], mask_1, contours[j], mask_2)
+
+            if any(collision_detected):
+                # print(f"Collision between {i} and {j} happen at step: {np.array(collision_detected).nonzero()}")
+                safety_actions[i] = 1  # Label collisions for OOIs now. Later we will build a larger dict.
+                safety_actions[j] = 1
+
+    # instead of loading a dict which saves all collision scenario, we could simply detect all agents' potential collision
+    return safety_actions
+
+
+def prepare_action_label(*, data_dict, dt, mask_probability, config):
+    """
+    mask_probability: the probability of masking the label. Should be around 0.05 or 0.1. Can't be too high.
+    """
+    ooi_ind = data_dict["decoder/labeled_agent_id"]
+    ooi_pos = utils.extract_data_by_agent_indices(data_dict["decoder/agent_position"], ooi_ind, agent_dim=1)[..., :2]
+    ooi_valid = utils.extract_data_by_agent_indices(
+        data_dict["decoder/agent_valid_mask"], ooi_ind, agent_dim=1
+    )  # (T, A)
+
+    # TODO: hardcoded here for now and we assume you can access GT trajectory. This won't work with test dataset.
+    assert ooi_pos.shape[0] == 91
+    assert ooi_valid.shape[0] == 91
+
+    # get the degree, acceleration, speed
+    turn_actions = get_direction_action_from_trajectory_batch(traj=ooi_pos, mask=ooi_valid, dt=dt, ooi=ooi_ind)
+    acce_actions = get_acce_action_from_trajectory_batch(ooi_pos, ooi_valid, dt=dt, ooi=ooi_ind)
+
+    # Rescatter labels to decoder-agent indices
+    assert config.TRAINING.PREDICT_ALL_AGENTS
+    B = data_dict["decoder/agent_valid_mask"].shape[1]
+
+    full_turn_actions = np.full((B, ), -1, dtype=int)
+    full_acce_actions = np.full((B, ), -1, dtype=int)
+
+    label_mask = np.random.binomial(1, mask_probability, size=len(ooi_ind))
+    label_invalid_mask = label_mask == 1
+
+    turn_actions[label_invalid_mask] = -1
+    acce_actions[label_invalid_mask] = -1
+
+    full_turn_actions[ooi_ind] = turn_actions
+    full_acce_actions[ooi_ind] = acce_actions
+
+    data_dict["decoder/label_turning"] = full_turn_actions
+    data_dict["decoder/label_acceleration"] = full_acce_actions
+
+    return data_dict
+
+
+def prepare_safety_label(*, data_dict, dt, mask_probability, config):
+    ooi_ind = data_dict["decoder/labeled_agent_id"]
+
+    ooi_pos = utils.extract_data_by_agent_indices(data_dict["decoder/agent_position"], ooi_ind, agent_dim=1)[..., :2]
+    ooi_valid = utils.extract_data_by_agent_indices(
+        data_dict["decoder/agent_valid_mask"], ooi_ind, agent_dim=1
+    )  # (T, A)
+
+    # TODO: hardcoded here for now and we assume you can access GT trajectory. This won't work with test dataset.
+    assert ooi_pos.shape[0] == 91
+    assert ooi_valid.shape[0] == 91
+
+    safety_actions = get_safety_action_from_trajectory_batch(data_dict, ooi_ind)
+
+    # Rescatter labels to decoder-agent indices
+    assert config.TRAINING.PREDICT_ALL_AGENTS
+    num_modeled_agents = data_dict["decoder/agent_valid_mask"].shape[1]
+
+    full_safety_actions = np.full((num_modeled_agents, ), -1, dtype=int)
+
+    label_mask = np.random.binomial(1, mask_probability, size=len(ooi_ind))
+    label_invalid_mask = label_mask == 1
+
+    label_invalid_mask[safety_actions == 1] = False  # We don't mask collision labels
+
+    safety_actions[label_invalid_mask] = -1
+
+    full_safety_actions[ooi_ind] = safety_actions
+
+    data_dict["decoder/label_safety"] = full_safety_actions
+
+    return data_dict
+
+
+if __name__ == '__main__':
+    scenario_dir = "/Users/claire_liu/validation_interactive_0/cat_reconstructed/sd_reconstructed_v0_ScenarioMap-21.pkl"
+    cat_dir = "/Users/claire_liu/validation_interactive_0/save.pkl"
+
+    import pickle
+
+    with open(scenario_dir, 'rb') as f:
+        scenario_data = pickle.load(f)
+    f.close()
+
+    with open(cat_dir, 'rb') as ff:
+        cat_dict = pickle.load(ff)
+    ff.close()
+
+    batch_labels = get_3d_action_label(scenario_data, cat_dict)
+    print(batch_labels)

scenestreamer/dataset/scenarionet_utils.py

@@ -0,0 +1,239 @@
+import numpy as np
+import torch
+from scenestreamer.utils import wrap_to_pi, rotate
+
+
+def overwrite_gt_to_pred_field(data_dict):
+    import copy
+    new_data_dict = copy.deepcopy(data_dict)
+    T, N, _ = data_dict["decoder/agent_position"].shape
+
+    new_data_dict["decoder/reconstructed_position"] = np.zeros((96, N, 2)).astype(np.float32)
+    new_data_dict["decoder/reconstructed_valid_mask"] = np.zeros((
+        96,
+        N,
+    )).astype(bool)
+    new_data_dict["decoder/reconstructed_heading"] = np.zeros((
+        96,
+        N,
+    )).astype(np.float32)
+    new_data_dict["decoder/reconstructed_velocity"] = np.zeros((96, N, 2)).astype(np.float32)
+
+    for id in range(N):  # overwrite all agents
+        traj = new_data_dict["decoder/agent_position"][:91, id, :2].astype(np.float32)
+        traj_mask = new_data_dict["decoder/agent_valid_mask"][:91, id].astype(bool)
+        theta = new_data_dict['decoder/agent_heading'][:91, id].astype(np.float32)
+        vel = new_data_dict['decoder/agent_velocity'][:91, id].astype(np.float32)
+
+        new_data_dict["decoder/reconstructed_position"][:91, id, :2] = traj
+        # new_data_dict["decoder/reconstructed_position"][:91, id, 2] = 0.0
+        new_data_dict["decoder/reconstructed_valid_mask"][:91, id] = traj_mask
+        # print(traj_mask)
+        new_data_dict["decoder/reconstructed_heading"][:91, id] = theta
+        new_data_dict["decoder/reconstructed_velocity"][:91, id] = vel
+
+    return new_data_dict
+
+
+def create_new_adv(data_dict):
+    ego_id = data_dict["decoder/sdc_index"]
+
+    ego_traj = data_dict["decoder/agent_position"][:, ego_id]
+    ego_heading = data_dict["decoder/agent_heading"][:, ego_id]
+    ego_velocity = data_dict["decoder/agent_velocity"][:, ego_id]
+    ego_shape = data_dict["decoder/agent_shape"][:, ego_id]
+    ego_mask = data_dict["decoder/agent_valid_mask"][:, ego_id]
+
+    last_valid_step = np.where(ego_mask)[0][-1]
+
+    # Create a new ADV at the final step.
+
+    adv_mask = np.zeros_like(ego_mask)
+    adv_mask[:last_valid_step + 1] = True
+
+    adv_traj = np.zeros_like(ego_traj)
+    adv_heading = np.zeros_like(ego_heading)
+    adv_velocity = np.zeros_like(ego_velocity)
+    adv_shape = np.zeros_like(ego_shape)
+
+    # Copy the final pos/head/vel/shape of ego
+    adv_traj[last_valid_step] = ego_traj[last_valid_step] + np.random.normal(loc=0.0, scale=0.5, size=3)
+    adv_heading[last_valid_step] = ego_heading[last_valid_step] + np.random.normal(loc=0.0, scale=0.1, size=1)
+    adv_velocity[last_valid_step] = ego_velocity[last_valid_step] + np.random.normal(loc=0.0, scale=0.5, size=2)
+
+    for i in range(data_dict["decoder/agent_shape"].shape[0]):
+        adv_shape[i] = ego_shape[last_valid_step]
+
+    # Insert data back:
+    data_dict["decoder/agent_position"] = np.concatenate(
+        [data_dict["decoder/agent_position"], adv_traj[:, None]], axis=1
+    )
+    data_dict["decoder/agent_heading"] = np.concatenate(
+        [data_dict["decoder/agent_heading"], adv_heading[:, None]], axis=1
+    )
+    data_dict["decoder/agent_velocity"] = np.concatenate(
+        [data_dict["decoder/agent_velocity"], adv_velocity[:, None]], axis=1
+    )
+    # data_dict["decoder/agent_shape"] = np.concatenate([data_dict["decoder/agent_shape"], adv_shape[:, None]], axis=1)
+
+    data_dict["decoder/agent_shape"] = np.concatenate([data_dict["decoder/agent_shape"], adv_shape[:, None]], axis=1)
+
+    data_dict["decoder/agent_valid_mask"] = np.concatenate(
+        [data_dict["decoder/agent_valid_mask"], adv_mask[:, None]], axis=1
+    )
+
+    data_dict["decoder/current_agent_shape"] = np.concatenate(
+        [data_dict["decoder/current_agent_shape"], data_dict["decoder/current_agent_shape"][ego_id:ego_id + 1]], axis=0
+    )
+    data_dict["decoder/agent_type"] = np.concatenate(
+        [data_dict["decoder/agent_type"], data_dict["decoder/agent_type"][ego_id:ego_id + 1]], axis=0
+    )
+    data_dict["decoder/agent_id"] = np.concatenate(
+        [data_dict["decoder/agent_id"], [len(data_dict["decoder/agent_id"])]], axis=0
+    )
+
+    # Add ADV into OOI:
+    data_dict["decoder/object_of_interest_id"] = np.concatenate(
+        [data_dict["decoder/object_of_interest_id"], [len(data_dict["decoder/agent_id"]) - 1]], axis=0
+    )
+
+    # Deal with some thing for forward prediction:
+    data_dict["decoder/current_agent_valid_mask"] = np.concatenate(
+        [data_dict["decoder/current_agent_valid_mask"], [1]], axis=0
+    )
+
+    print("====================================")
+    print(
+        "The new ADV is created at the final step {}, it's ID is: {}".format(
+            last_valid_step,
+            len(data_dict["decoder/agent_id"]) - 1
+        )
+    )
+    print("====================================")
+
+    return data_dict
+
+
+def overwrite_to_scenario_description(output_dict_mode, original_SD, ooi=None, adv_id=None):
+    # overwrite original SD with all predicted ooi trajectories included
+    # import pdb; pdb.set_trace()
+    if not ooi:
+        ooi = output_dict_mode['decoder/agent_id']  # overwrite all agents
+    sdc_track_name = original_SD['metadata']['sdc_id']
+    adv_track_name = str(output_dict_mode['decoder/track_name'][int(adv_id)].item())
+
+    for id in ooi:
+        agent_track_name = str(output_dict_mode['decoder/track_name'][id].item())
+
+        # begin to overwrite original scenario_data
+        agent_traj = output_dict_mode["decoder/agent_position"][:91, id, ]
+        agent_heading = output_dict_mode["decoder/agent_heading"][:91, id]
+        agent_vel = output_dict_mode["decoder/agent_velocity"][:91, id]
+        agent_traj_mask = output_dict_mode["decoder/agent_valid_mask"][:91, id]
+
+        # modify adv info
+        # agent_z = original_SD['tracks'][agent_track_name]['state']['position'][10, 2]  # fill the z-axis
+        # agent_traj_z = np.full((91, 1), agent_z)
+        # agent_new_traj = np.concatenate([agent_traj, agent_traj_z], axis=1)
+        # print("new_traj:", agent_new_traj.shape)
+        original_SD['tracks'][agent_track_name]['state']['position'] = agent_traj
+        original_SD['tracks'][agent_track_name]['state']['velocity'] = agent_vel
+        original_SD['tracks'][agent_track_name]['state']['heading'] = agent_heading
+        original_SD['tracks'][agent_track_name]['state']['valid'] = agent_traj_mask
+
+        length = original_SD['tracks'][agent_track_name]['state']['length'][10]
+        width = original_SD['tracks'][agent_track_name]['state']['width'][10]
+        height = original_SD['tracks'][agent_track_name]['state']['height'][10]
+        original_SD['tracks'][agent_track_name]['state']['length'] = np.full((91, ), length)
+        original_SD['tracks'][agent_track_name]['state']['width'] = np.full((91, ), width)
+        original_SD['tracks'][agent_track_name]['state']['height'] = np.full((91, ), height)
+
+    original_SD['metadata']['selected_adv_id'] = adv_track_name
+
+    return original_SD
+
+
+def overwrite_to_scenario_description_new_agent(output_dict_mode, original_SD, ooi=None):
+    # overwrite original SD with all predicted ooi trajectories included
+    ooi = output_dict_mode['decoder/agent_id']  # overwrite all agents
+
+    adv_track_name = 'new_adv_agent'
+    original_SD['tracks'][adv_track_name] = {'state': {}, 'type': 'VEHICLE', 'metadata': {}}
+    sdc_track_name = original_SD['metadata']['sdc_id']
+
+    for id in ooi:
+        if id == ooi[-1]:
+            agent_track_name = 'new_adv_agent'
+        else:
+            agent_track_name = str(output_dict_mode['decoder/track_name'][id].item())
+
+        # begin to overwrite original scenario_data
+        agent_traj = output_dict_mode["decoder/agent_position"][:, id, ]
+        agent_heading = output_dict_mode["decoder/agent_heading"][:, id]
+        agent_vel = output_dict_mode["decoder/agent_velocity"][:, id]
+        agent_traj_mask = output_dict_mode["decoder/agent_valid_mask"][:, id]
+
+        # modify adv info
+        # agent_z = original_SD['tracks'][agent_track_name]['state']['position'][10, 2]  # fill the z-axis
+        # agent_traj_z = np.full((91, 1), agent_z)
+        # agent_new_traj = np.concatenate([agent_traj, agent_traj_z], axis=1)
+        # print("new_traj:", agent_new_traj.shape)
+        original_SD['tracks'][agent_track_name]['state']['position'] = agent_traj
+
+        original_SD['tracks'][agent_track_name]['state']['velocity'] = agent_vel
+        original_SD['tracks'][agent_track_name]['state']['heading'] = agent_heading
+        original_SD['tracks'][agent_track_name]['state']['valid'] = agent_traj_mask
+
+        length = original_SD['tracks'][sdc_track_name]['state']['length'][10]
+        width = original_SD['tracks'][sdc_track_name]['state']['width'][10]
+        height = original_SD['tracks'][sdc_track_name]['state']['height'][10]
+        original_SD['tracks'][agent_track_name]['state']['length'] = np.full((91, ), length)
+        original_SD['tracks'][agent_track_name]['state']['width'] = np.full((91, ), width)
+        original_SD['tracks'][agent_track_name]['state']['height'] = np.full((91, ), height)
+
+    original_SD['tracks'][adv_track_name]['metadata']['dataset'] = 'waymo'
+    original_SD['tracks'][adv_track_name]['metadata']['object_id'] = 'new_adv_agent'
+    original_SD['tracks'][adv_track_name]['metadata']['track_length'] = 91
+    original_SD['tracks'][adv_track_name]['metadata']['type'] = 'VEHICLE'
+    original_SD['metadata']['new_adv_id'] = 'new_adv_agent'
+    original_SD['metadata']['objects_of_interest'].append('new_adv_agent')
+    tracks_length = len(list(original_SD['tracks'].keys()))
+    original_SD['metadata']['tracks_to_predict']['new_adv_agent'] = {
+        'difficulty': 0,
+        'object_type': 'VEHICLE',
+        'track_id': 'new_adv_agent',
+        'track_index': tracks_length - 1
+    }
+
+    return original_SD
+
+
+def transform_to_global_coordinate(data_dict):
+    map_center = data_dict["metadata/map_center"].reshape(-1, 1, 3)  # (1,1,3)
+    assert "decoder/agent_position" in data_dict, "Have you set EVALUATION.PREDICT_ALL_AGENTS to False?"
+    T, N, _ = data_dict["decoder/agent_position"].shape
+    assert data_dict["decoder/agent_position"].ndim == 3
+    data_dict["decoder/agent_position"] += map_center
+
+    return data_dict
+
+
+def _overwrite_datadict_all_agents(data_dict):
+    import copy
+    new_data_dict = copy.deepcopy(data_dict)
+
+    T, N, _ = data_dict["decoder/reconstructed_position"].shape
+
+    for id in range(N):  # overwrite all agents
+        traj = data_dict["decoder/reconstructed_position"][:91, id, ]
+        traj_mask = data_dict["decoder/reconstructed_valid_mask"][:91, id]
+        theta = data_dict['decoder/reconstructed_heading'][:91, id]
+        vel = data_dict['decoder/reconstructed_velocity'][:91, id]
+
+        new_data_dict["decoder/agent_position"][:, id, :2] = traj
+        new_data_dict["decoder/agent_position"][:, id, 2] = 0.0
+        new_data_dict["decoder/agent_valid_mask"][:, id] = traj_mask
+        new_data_dict["decoder/agent_heading"][:, id] = theta
+        new_data_dict["decoder/agent_velocity"][:, id] = vel
+
+    return new_data_dict