ameliacmu/ameliatf

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.