# ameliacmu/ameliatf architecture

> Architecture diagram of the GitHub repository ameliacmu/ameliatf, generated by GitDiagram from its file tree, README and a sample of source files (last updated 2026-06-26). It is an AI-made overview: check details against the source.

- Interactive diagram: https://gitdiagram.com/ameliacmu/ameliatf
- Repository: https://github.com/ameliacmu/ameliatf

## Overview

1. **What this repo is**
   - `AmeliaTF` is the training/evaluation code for an airport-surface trajectory forecasting model, not the raw dataset itself.
   - The main external dependency is the **Amelia-48** dataset, which is expected to be preprocessed by the companion AmeliaScenes tooling and linked into `datasets/`.

2. **Top-level execution flow**
   - The two runtime entrypoints are `amelia_tf/train.py` and `amelia_tf/eval.py`.
   - Both are driven by **Hydra** configuration, so architecture is assembled from config groups rather than hard-coded scripts.

3. **Configuration as the composition layer**
   - `configs/` defines the experiment graph: data splits, model variants, trainer backends, callbacks, loggers, and debug modes.
   - The repo is organized around reusable experiment presets such as `single-airport`, `seen-*`, `joint`, `marginal`, and trajectory-only variants.
   - Trainer configs cover CPU/GPU/DDP/MPS paths, which matters because the README notes stable runs are currently multi-GPU-oriented.

4. **Data subsystem**
   - `amelia_tf/data/components/` contains the dataset primitives:
     - `base_dataset.py` for shared dataset behavior,
     - `amelia_dataset.py` for Amelia-specific scene samples.
   - `datamodule.py` and `datamodule_from_list.py` wrap those datasets into **PyTorch Lightning DataModules**, which are the boundary between raw sample loading and model consumption.
   - The data configs under `configs/data/` encode airport-specific and multi-airport splits, so the dataset layer is where “seen vs unseen airport” experiment boundaries are enforced.

5. **Model subsystem**
   - `amelia_tf/models/trajpred.py` appears to be the main Lightning model wrapper for trajectory prediction.
   - `amelia_tf/models/components/` contains the actual architectural blocks:
     - `amelia.py` and `amelia_traj.py` for the main model variants,
     - `context.py` for map/context encoding,
     - `self_attention.py` and `cross_attention.py` for transformer-style scene encoding,
     - `gmm.py` for the probabilistic trajectory decoder,
     - `common.py` for shared layers/utilities.
   - The README’s described pipeline is: scene representation/scoring → hierarchical transformer scene encoding → GMM-based future prediction.

6. **Model variants and boundaries**
   - The repository supports both **marginal** and **joint** prediction paradigms.
   - It also supports **context-aware** and **trajectory-only** versions via `marginal_traj` / `joint_traj`, which likely bypass the map/context branch while keeping the forecasting stack intact.
   - That split is an important diagram boundary: shared forecasting core, optional context encoder.

7. **Training loop and experiment management**
   - Training is built on **PyTorch Lightning**, with callbacks configured in `configs/callbacks/` for checkpointing, early stopping, LR monitoring, model summaries, and progress bars.
   - `amelia_tf/utils/instantiators.py` and related helpers likely glue Hydra-configured objects into Lightning modules, datamodules, loggers, and callbacks.
   - This is a config-driven experiment system rather than a bespoke training loop.

8. **Evaluation flow**
   - `amelia_tf/eval.py` is a standalone inference/evaluation path that loads a checkpoint and computes forecasting metrics.
   - The README indicates this is also config-driven and expects an explicit `ckpt_path`.
   - Evaluation is therefore a separate pipeline boundary from training, sharing the same data/model configs but not optimizer state.

9. **Metrics, losses, and masking**
   - `amelia_tf/utils/losses.py`, `metrics.py`, and `global_masks.py` handle the forecasting objective and metric computation.
   - The presence of global masks suggests scene-level validity filtering and masked aggregation across agents/timesteps.
   - This is likely where ADE/FDE-style reporting is implemented, matching the paper’s metrics.

10. **Data quality and preprocessing support**
    - `datasets/error_files.py` suggests explicit handling of known problematic samples or preprocessing failures.
    - The repo depends on prior dataset post-processing, so the architecture assumes clean, structured scenario files rather than raw sensor feeds.

11. **Logging and observability**
    - `configs/logger/` supports multiple experiment tracking backends: TensorBoard, WandB, MLflow, Neptune, Comet, Aim, and CSV.
    - This makes logging a plug-in boundary rather than a core dependency.

12. **Repository/runtime conventions**
    - The project follows the **lightning-hydra-template** pattern, with `.hydra/` storing resolved run configs and `pyproject.toml` defining Python package metadata.
    - In architecture terms, this is a packaged research codebase where configuration, not code branching, selects most runtime behavior.

13. **Main data path to diagram**
    - External Amelia dataset → processed scenario files in `datasets/` → `AmeliaDataset`/`DataModule` → Lightning model (`trajpred` + component blocks) → loss/metrics → logger/checkpoints.
    - Evaluation reuses the same dataset/model path but starts from a saved checkpoint instead of optimizer training.

14. **Key diagram boundaries to preserve**
    - Separate **external data preprocessing** from this repo’s training code.
    - Separate **configuration assembly** from execution.
    - Separate **data loading**, **model computation**, and **experiment tracking**.
    - Show optional branches for **context vs trajectory-only** and **marginal vs joint** prediction, since those are first-class architectural variations here.

Main entry points: [`configs`](https://github.com/ameliacmu/ameliatf/tree/main/configs) (Hydra cfg), [`amelia_tf/models/trajpred.py`](https://github.com/ameliacmu/ameliatf/blob/main/amelia_tf/models/trajpred.py) (TrajPred), [`amelia_tf/utils/instantiators.py`](https://github.com/ameliacmu/ameliatf/blob/main/amelia_tf/utils/instantiators.py) (Instantiators), [`amelia_tf/models/components/amelia.py`](https://github.com/ameliacmu/ameliatf/blob/main/amelia_tf/models/components/amelia.py) (Amelia core)

## Diagram (Mermaid)

```mermaid
flowchart TD

subgraph group_external["External data"]
  node_amelia48[("Amelia-48<br/>external dataset")]
  node_datasets_dir[("Datasets<br/>linked data")]
end

subgraph group_config["Configuration"]
  node_hydra_cfg["Hydra cfg<br/>composition layer"]
  node_data_cfg["Data presets<br/>split config"]
  node_model_cfg["Model presets<br/>variant config"]
  node_trainer_cfg["Trainer presets<br/>runtime config"]
  node_logger_cfg["Logger presets<br/>tracking config"]
  node_callback_cfg["Callbacks<br/>training config"]
end

subgraph group_runtime["Execution"]
  node_train_entry["Train<br/>entrypoint<br/>[train.py]"]
  node_eval_entry["Eval<br/>entrypoint<br/>[eval.py]"]
end

subgraph group_data["Data subsystem"]
  node_amelia_dataset["AmeliaDataset<br/>[amelia_dataset.py]"]
  node_base_dataset["BaseDataset<br/>dataset base<br/>[base_dataset.py]"]
  node_datamodule["DataModule<br/>[datamodule.py]"]
  node_datamodule_list["List DataModule<br/>lightning datamodule"]
end

subgraph group_model["Model subsystem"]
  node_trajpred["TrajPred<br/>lightning module<br/>[trajpred.py]"]
  node_amelia_model["Amelia core<br/>forecasting core<br/>[amelia.py]"]
  node_amelia_traj["Traj core<br/>forecasting core<br/>[amelia_traj.py]"]
  node_context["Context encoder<br/>scene encoder<br/>[context.py]"]
  node_attn_blocks["Attention blocks<br/>transformer blocks<br/>[self_attention.py]"]
  node_gmm["GMM decoder<br/>probabilistic decoder<br/>[gmm.py]"]
end

subgraph group_utils["Utilities"]
  node_losses_metrics["Losses & metrics<br/>objective and eval<br/>[losses.py]"]
  node_instantiators["Instantiators<br/>wiring helpers<br/>[instantiators.py]"]
  node_mask_utils["Global masks<br/>masking utils<br/>[global_masks.py]"]
  node_log_utils["Logging utils<br/>observability<br/>[logging_utils.py]"]
end

node_train_entry -->|"loads"| node_hydra_cfg
node_eval_entry -->|"loads"| node_hydra_cfg
node_hydra_cfg -->|"includes"| node_data_cfg
node_hydra_cfg -->|"includes"| node_model_cfg
node_hydra_cfg -->|"includes"| node_trainer_cfg
node_hydra_cfg -->|"includes"| node_logger_cfg
node_hydra_cfg -->|"includes"| node_callback_cfg
node_data_cfg -->|"points to"| node_datasets_dir
node_datasets_dir -->|"derived from"| node_amelia48
node_train_entry -->|"wires"| node_instantiators
node_eval_entry -->|"wires"| node_instantiators
node_instantiators -->|"creates"| node_datamodule
node_instantiators -->|"creates"| node_datamodule_list
node_datamodule -->|"wraps"| node_amelia_dataset
node_datamodule_list -->|"wraps"| node_amelia_dataset
node_amelia_dataset -->|"extends"| node_base_dataset
node_instantiators -->|"creates"| node_trajpred
node_trajpred -->|"uses"| node_amelia_model
node_trajpred -->|"variant"| node_amelia_traj
node_amelia_model -->|"includes"| node_context
node_amelia_model -->|"includes"| node_attn_blocks
node_amelia_traj -->|"includes"| node_attn_blocks
node_amelia_model -->|"decodes with"| node_gmm
node_amelia_traj -->|"decodes with"| node_gmm
node_trajpred -->|"optimizes"| node_losses_metrics
node_trajpred -->|"filters"| node_mask_utils
node_train_entry -->|"logs via"| node_log_utils
node_eval_entry -->|"logs via"| node_log_utils
node_logger_cfg -->|"configures"| node_log_utils
node_callback_cfg -->|"monitors"| node_trajpred

click node_datasets_dir "https://github.com/ameliacmu/ameliatf/tree/main/datasets"
click node_train_entry "https://github.com/ameliacmu/ameliatf/blob/main/amelia_tf/train.py"
click node_eval_entry "https://github.com/ameliacmu/ameliatf/blob/main/amelia_tf/eval.py"
click node_hydra_cfg "https://github.com/ameliacmu/ameliatf/tree/main/configs"
click node_data_cfg "https://github.com/ameliacmu/ameliatf/tree/main/configs/data"
click node_model_cfg "https://github.com/ameliacmu/ameliatf/tree/main/configs/model"
click node_trainer_cfg "https://github.com/ameliacmu/ameliatf/tree/main/configs/trainer"
click node_logger_cfg "https://github.com/ameliacmu/ameliatf/tree/main/configs/logger"
click node_callback_cfg "https://github.com/ameliacmu/ameliatf/tree/main/configs/callbacks"
click node_amelia_dataset "https://github.com/ameliacmu/ameliatf/blob/main/amelia_tf/data/components/amelia_dataset.py"
click node_base_dataset "https://github.com/ameliacmu/ameliatf/blob/main/amelia_tf/data/components/base_dataset.py"
click node_datamodule "https://github.com/ameliacmu/ameliatf/blob/main/amelia_tf/data/datamodule.py"
click node_datamodule_list "https://github.com/ameliacmu/ameliatf/blob/main/amelia_tf/data/datamodule_from_list.py"
click node_trajpred "https://github.com/ameliacmu/ameliatf/blob/main/amelia_tf/models/trajpred.py"
click node_amelia_model "https://github.com/ameliacmu/ameliatf/blob/main/amelia_tf/models/components/amelia.py"
click node_amelia_traj "https://github.com/ameliacmu/ameliatf/blob/main/amelia_tf/models/components/amelia_traj.py"
click node_context "https://github.com/ameliacmu/ameliatf/blob/main/amelia_tf/models/components/context.py"
click node_attn_blocks "https://github.com/ameliacmu/ameliatf/blob/main/amelia_tf/models/components/self_attention.py"
click node_gmm "https://github.com/ameliacmu/ameliatf/blob/main/amelia_tf/models/components/gmm.py"
click node_losses_metrics "https://github.com/ameliacmu/ameliatf/blob/main/amelia_tf/utils/losses.py"
click node_instantiators "https://github.com/ameliacmu/ameliatf/blob/main/amelia_tf/utils/instantiators.py"
click node_mask_utils "https://github.com/ameliacmu/ameliatf/blob/main/amelia_tf/utils/global_masks.py"
click node_log_utils "https://github.com/ameliacmu/ameliatf/blob/main/amelia_tf/utils/logging_utils.py"

classDef toneNeutral fill:#f8fafc,stroke:#334155,stroke-width:1.5px,color:#0f172a
classDef toneBlue fill:#dbeafe,stroke:#2563eb,stroke-width:1.5px,color:#172554
classDef toneAmber fill:#fef3c7,stroke:#d97706,stroke-width:1.5px,color:#78350f
classDef toneMint fill:#dcfce7,stroke:#16a34a,stroke-width:1.5px,color:#14532d
classDef toneRose fill:#ffe4e6,stroke:#e11d48,stroke-width:1.5px,color:#881337
classDef toneIndigo fill:#e0e7ff,stroke:#4f46e5,stroke-width:1.5px,color:#312e81
classDef toneTeal fill:#ccfbf1,stroke:#0f766e,stroke-width:1.5px,color:#134e4a
class node_amelia48,node_datasets_dir toneBlue
class node_hydra_cfg,node_data_cfg,node_model_cfg,node_trainer_cfg,node_logger_cfg,node_callback_cfg toneAmber
class node_train_entry,node_eval_entry toneMint
class node_amelia_dataset,node_base_dataset,node_datamodule,node_datamodule_list toneRose
class node_trajpred,node_amelia_model,node_amelia_traj,node_context,node_attn_blocks,node_gmm toneIndigo
class node_losses_metrics,node_instantiators,node_mask_utils,node_log_utils toneTeal
```

## Components

### External data

Preprocessed Amelia-48 scenarios linked into the repo

- **Amelia-48**: external dataset · Preprocessed airport-surface trajectory data consumed by this repo
- [Datasets](https://github.com/ameliacmu/ameliatf/tree/main/datasets): linked data · Local scenario files and processed dataset inputs · `datasets`

### Configuration

Hydra experiment composition and runtime presets

- [Hydra cfg](https://github.com/ameliacmu/ameliatf/tree/main/configs): composition layer · Top-level config assembly for runs · `configs`
- [Data presets](https://github.com/ameliacmu/ameliatf/tree/main/configs/data): split config · Airport-specific and multi-airport data splits · `configs/data`
- [Model presets](https://github.com/ameliacmu/ameliatf/tree/main/configs/model): variant config · Marginal, joint, and trajectory-only model variants · `configs/model`
- [Trainer presets](https://github.com/ameliacmu/ameliatf/tree/main/configs/trainer): runtime config · CPU, GPU, DDP, and MPS training backends · `configs/trainer`
- [Logger presets](https://github.com/ameliacmu/ameliatf/tree/main/configs/logger): tracking config · Experiment logging backend selection · `configs/logger`
- [Callbacks](https://github.com/ameliacmu/ameliatf/tree/main/configs/callbacks): training config · Checkpointing, early stopping, progress, and monitoring · `configs/callbacks`

### Execution

Training and evaluation entrypoints

- [Train](https://github.com/ameliacmu/ameliatf/blob/main/amelia_tf/train.py): entrypoint · Hydra-driven training script · `amelia_tf/train.py`
- [Eval](https://github.com/ameliacmu/ameliatf/blob/main/amelia_tf/eval.py): entrypoint · Hydra-driven evaluation script · `amelia_tf/eval.py`

### Data subsystem

Dataset and datamodule boundaries

- [AmeliaDataset](https://github.com/ameliacmu/ameliatf/blob/main/amelia_tf/data/components/amelia_dataset.py): dataset · Amelia-specific scene sample loading and formatting · `amelia_tf/data/components/amelia_dataset.py`
- [BaseDataset](https://github.com/ameliacmu/ameliatf/blob/main/amelia_tf/data/components/base_dataset.py): dataset base · Shared dataset behavior · `amelia_tf/data/components/base_dataset.py`
- [DataModule](https://github.com/ameliacmu/ameliatf/blob/main/amelia_tf/data/datamodule.py): lightning datamodule · PyTorch Lightning boundary between loading and model input · `amelia_tf/data/datamodule.py`
- [List DataModule](https://github.com/ameliacmu/ameliatf/blob/main/amelia_tf/data/datamodule_from_list.py): lightning datamodule · DataModule variant built from explicit file lists · `amelia_tf/data/datamodule_from_list.py`

### Model subsystem

Forecasting wrapper and internal components

- [TrajPred](https://github.com/ameliacmu/ameliatf/blob/main/amelia_tf/models/trajpred.py): lightning module · Main forecasting wrapper around the Amelia model · `amelia_tf/models/trajpred.py`
- [Amelia core](https://github.com/ameliacmu/ameliatf/blob/main/amelia_tf/models/components/amelia.py): forecasting core · Main marginal/joint prediction model · `amelia_tf/models/components/amelia.py`
- [Traj core](https://github.com/ameliacmu/ameliatf/blob/main/amelia_tf/models/components/amelia_traj.py): forecasting core · Trajectory-only variant without full context branch · `amelia_tf/models/components/amelia_traj.py`
- [Context encoder](https://github.com/ameliacmu/ameliatf/blob/main/amelia_tf/models/components/context.py): scene encoder · Map and airport context encoding · `amelia_tf/models/components/context.py`
- [Attention blocks](https://github.com/ameliacmu/ameliatf/blob/main/amelia_tf/models/components/self_attention.py): transformer blocks · Self-attention and cross-attention scene encoding · `amelia_tf/models/components/self_attention.py`
- [GMM decoder](https://github.com/ameliacmu/ameliatf/blob/main/amelia_tf/models/components/gmm.py): probabilistic decoder · Gaussian mixture future trajectory prediction · `amelia_tf/models/components/gmm.py`

### Utilities

Losses, metrics, logging, and instantiation helpers

- [Losses & metrics](https://github.com/ameliacmu/ameliatf/blob/main/amelia_tf/utils/losses.py): objective and eval · Forecasting loss, ADE/FDE-style metrics, and masking · `amelia_tf/utils/losses.py`
- [Instantiators](https://github.com/ameliacmu/ameliatf/blob/main/amelia_tf/utils/instantiators.py): wiring helpers · Hydra-to-Lightning object construction helpers · `amelia_tf/utils/instantiators.py`
- [Global masks](https://github.com/ameliacmu/ameliatf/blob/main/amelia_tf/utils/global_masks.py): masking utils · Scene-level validity filtering and aggregation masks · `amelia_tf/utils/global_masks.py`
- [Logging utils](https://github.com/ameliacmu/ameliatf/blob/main/amelia_tf/utils/logging_utils.py): observability · Logger setup and rich logging helpers · `amelia_tf/utils/logging_utils.py`

## Connections

- Train → Hydra cfg: loads · Training is assembled from Hydra configs
- Eval → Hydra cfg: loads · Evaluation is assembled from Hydra configs
- Hydra cfg → Data presets: includes · Selects dataset split presets
- Hydra cfg → Model presets: includes · Selects model variant presets
- Hydra cfg → Trainer presets: includes · Selects execution backend preset
- Hydra cfg → Logger presets: includes · Selects experiment tracking backend
- Hydra cfg → Callbacks: includes · Selects training callbacks
- Data presets → Datasets: points to · Experiment splits resolve into linked dataset files
- Datasets → Amelia-48: derived from · Local datasets are processed Amelia-48 scenarios
- Train → Instantiators: wires · Builds runtime objects from config
- Eval → Instantiators: wires · Builds runtime objects from config
- Instantiators → DataModule: creates · Constructs the main Lightning DataModule
- Instantiators → List DataModule: creates · Constructs list-based data loading when needed
- DataModule → AmeliaDataset: wraps · Exposes Amelia samples to the trainer
- List DataModule → AmeliaDataset: wraps · Builds datasets from explicit file lists
- AmeliaDataset → BaseDataset: extends · Shares common dataset behavior
- Instantiators → TrajPred: creates · Constructs the Lightning forecasting module
- TrajPred → Amelia core: uses · Delegates core forecasting to the Amelia model
- TrajPred → Traj core: variant · Switches to trajectory-only forecasting core
- Amelia core → Context encoder: includes · Optional context-aware scene encoding branch
- Amelia core → Attention blocks: includes · Uses self/cross-attention for hierarchical encoding
- Traj core → Attention blocks: includes · Retains forecasting encoder blocks
- Amelia core → GMM decoder: decodes with · Produces probabilistic future trajectories
- Traj core → GMM decoder: decodes with · Produces probabilistic future trajectories
- TrajPred → Losses & metrics: optimizes · Training objective and reporting metrics
- TrajPred → Global masks: filters · Applies scene/agent/timestep validity masks
- Train → Logging utils: logs via · Training emits experiment logs and summaries
- Eval → Logging utils: logs via · Evaluation emits metrics and run output
- Logger presets → Logging utils: configures · Selects the active logging backend
- Callbacks → TrajPred: monitors · Callbacks observe checkpoints and metrics during training