CITRIS-joint-embedding-Voice2Traj

Python package for joint embedding of ATC voice communications and aircraft trajectories. The goal is to learn a shared representation space where a voice segment and its corresponding trajectory sit close together, enabling cross-modal retrieval (voice → next trajectory, trajectory → matching voice).

The project is part of the CITRIS aviation research context and is designed to be continued — each training stage is independent, checkpointed, and documented with a PDF report. See pdf_latex_reports/ for results so far.

Table of Contents

1. Pipeline Overview
2. Installation
3. Data Setup
4. Repository Structure
5. Training Stages
   • 1. Voice Feature Extraction
   • 2. Trajectory Feature Extraction
   • 3. Contrastive Learning
   • 4. Bijective Mapping (RealNVP)
   • 5. Next Step Prediction
   • 6. Next Sentence Prediction
   • 7. Ablation Study
6. Visualization & Notebooks
7. Public API & Models
8. Testing
9. PDF Reports
10. Sources and References

---

Pipeline Overview

Raw audio (.mp3)                    Raw trajectory (.nc / OpenSky)
      │                                         │
      ▼                                         ▼
 ATCTranscriptor                     MAE Encoder (external repo)
 (Whisper or Wav2Vec2)               CITRIS-modern-transformer-encoding
      │                                         │
      ▼                                         ▼
Voice embeddings (.h5)              Trajectory embeddings (.h5)
  dim=1280 (Whisper)                  dim=1792 (14 × 128)
  dim=1024 (Wav2Vec2)
      │                                         │
      └─────────────┬───────────────────────────┘
                    ▼
         Contrastive learning (InfoNCE)
         FFNN projectors → joint space (dim=2048)
                    │
                    ▼
         Bijective mapping (RealNVP)
         joint ↔ voice / joint ↔ traj
                    │
          ┌─────────┴─────────┐
          ▼                   ▼
   kNN lookup table     Next-step predictor
   (voice → traj)       (partial traj → future)

Each stage saves checkpoints under checkpoints/ and results under results/.

---

Installation

git clone https://github.com/LouisBrusset/CITRIS-joint-embedding-Voice2Traj.git
cd CITRIS-joint-embedding-Voice2Traj

Create a Python 3.12 virtual environment (required for the CUDA torch wheels):

# With uv (recommended)
uv venv --python 3.12
.venv\Scripts\Activate.ps1

# Or plain venv
python -m venv .venv
.venv\Scripts\Activate.ps1

Install in editable mode:

uv pip install -e .

ffmpeg must be in PATH for audio processing (.mp3 input):

# Windows — with winget or choco
winget install ffmpeg

HuggingFace token (only needed if re-downloading models from the Hub):

• Set HF_TOKEN in the environment, or
• Place a huggingface_token.json at the repo root or next to src/joint_embedding_voice2traj/model/wav2vect.py.

Local copies of the pretrained models are already cached under models/.

---

Data Setup

The training scripts expect data on a separate drive. Before running any training, update the path constants at the top of the relevant config file.

The dataset contains recordings from 8 ATC towers:

knuq_gnd_twr  koak_del_gnd_twr  koak_norcal_app
ksfo_clnc_del  ksfo_gnd_twr  ksfo_norcal_app
ksjc_norcal_app  ksjc_twr

Trajectory embeddings come from the MAE encoder trained in the companion repo CITRIS-modern-transformer-encoding. The config_next_step.py points to that checkpoint:

MAE_CHECKPOINT = "../CITRIS-modern-transformer-encoding/checkpoints/stage1_transformer_opensky_v0/best.pt"

---

Repository Structure

CITRIS-joint-embedding-Voice2Traj/
├── checkpoints/                  # Saved model weights (gitignored)
│   ├── ablation/
│   ├── bijective/
│   ├── contrastive/
│   ├── next_sentence_prediction/
│   └── next_step/
├── data/
│   └── step_1_trajectories/      # Raw trajectory data
├── dev/
│   └── test.ps1                  # Windows test runner
├── models/                       # Cached pretrained models
│   ├── wav2vec2-xls-r-300m-en-atc-uwb-atcc-and-atcosim/
│   ├── whisper-large-v3/
│   └── models--Systran--faster-whisper-large-v3/
├── notebooks/                    # Jupyter notebooks
├── pdf_latex_reports/            # LaTeX-generated PDF reports per stage
├── results/                      # Plots and metrics per training run
├── src/
│   ├── joint_embedding_voice2traj/
│   │   ├── __init__.py           # Public API (lazy imports)
│   │   ├── model/                # Core model implementations
│   │   ├── ml_tools/             # Device selection, seeds, callbacks
│   │   ├── features_projection/  # UMAP visualization
│   │   └── training/             # One sub-package per training stage
│   │       ├── voice_features_inference/
│   │       ├── traj_features_inference/
│   │       ├── contrastive_learning/
│   │       ├── bijective_learning/
│   │       ├── next_step_prediction/
│   │       ├── next_sentence_prediction/
│   │       └── ablation_study/
│   └── tests/
├── Makefile
├── pyproject.toml
└── uv.lock

---

Training Stages

Each stage has a config_*.py file that centralizes all hyperparameters. Edit that file, then run the corresponding training.py (or run_all.py for the ablation).

1. Voice Feature Extraction

Config: src/joint_embedding_voice2traj/training/voice_features_inference/config_voice.py

Extracts fixed-size voice embeddings from raw audio files and saves them to H5 files.

Output dimensions:

• Whisper large-v3: 1280
• Wav2Vec2 XLS-R: 1024

Output files: {tower}_embeddings_{model}_{pool_tag}.h5 in OUTPUT_DIR.

2. Trajectory Feature Extraction

Script: src/joint_embedding_voice2traj/training/traj_features_inference/retrieve_traj_features.py

Runs the MAE encoder (from CITRIS-modern-transformer-encoding) over trajectory sequences and saves the per-timestep hidden states. Output dimension: 14 timesteps × 128 d_model = 1792 (flattened).

3. Contrastive Learning

Config: src/joint_embedding_voice2traj/training/contrastive_learning/config_contrastive.py

Trains two FFNN projectors (one for voice, one for trajectory) that push paired embeddings into a shared joint space using the InfoNCE (contrastive) loss.

Checkpoints → checkpoints/contrastive/. A sweep variant (sweep_training.py) runs a grid over projector architectures.

4. Bijective Mapping (RealNVP)

Config: src/joint_embedding_voice2traj/training/bijective_learning/config_bijective.py

Trains a RealNVP normalizing flow that maps bidirectionally between the joint space and each modality space. This makes the joint embedding fully invertible: given a joint vector, you can reconstruct a voice or trajectory embedding.

The training set is augmented with synthetic data to improve coverage of the joint space:

RealNVP architecture:

Loss weights:

Training:

Checkpoints → checkpoints/bijective/.

5. Next Step Prediction

Config: src/joint_embedding_voice2traj/training/next_step_prediction/config_next_step.py

Trains a lightweight MLP (StepPredictor) on top of the frozen MAE encoder. Given the encoder hidden states for the first 10 timesteps, it predicts the next 4 timesteps.

Sequence layout:

StepPredictor MLP: (B, 4, 128) → flatten → (B, 512) → [512, 256] → (B, 4, 6)

Training:

Checkpoints → checkpoints/next_step/. The contrastive and bijective checkpoints are loaded for evaluation only (not trained here).

6. Next Sentence Prediction

Config: src/joint_embedding_voice2traj/training/next_sentence_prediction/config_nsp.py

Cross-modal retrieval: given a voice embedding, find the most similar trajectory in the lookup table.

Pipeline modes — controlled by USE_FFNN_PROJ:

kNN parameters:

The lookup table is precomputed from the training set and saved to checkpoints/next_sentence_prediction/lookup_table.npz. Run build_lookup_table.py once before using pipeline.py.

7. Ablation Study

Config: src/joint_embedding_voice2traj/training/ablation_study/config_ablation.py

Compares two end-to-end architectures:

Both share the same RealNVP config: 8 blocks, hidden [256, 256], ActNorm=True, Tanh=True.

Run: python src/joint_embedding_voice2traj/training/ablation_study/run_all.py

---

Visualization & Notebooks

UMAP projections

Config: src/joint_embedding_voice2traj/features_projection/config_proj.py

Notebooks

---

Public API & Models

Once installed, all main classes are importable from the package root:

from joint_embedding_voice2traj import (
    # Contrastive
    FFNN, InfoNCELoss, PairDataset,
    # Bijective
    RealNVP1D, AffineCoupling1D, ActNorm1D, ConditionerMLP,
    BijectiveReconstructionLoss, GaussianPriorLoss,
    # ATC speech
    ATCEncoder, ATCDecoder, ATCTranscriptor,
    # Trajectory
    LSTMPredictor, TimeSeriesDataset, pad_collate,
    xLSTM, sLSTMCell, mLSTMCell, xLSTMBlock, TrajectoryMSELoss,
)

You can also import from submodules directly:

from joint_embedding_voice2traj.model.real_nvp import RealNVP1D
from joint_embedding_voice2traj.model.x_lstm import xLSTM

Quick example: contrastive training step

import torch
from joint_embedding_voice2traj import FFNN, InfoNCELoss

voice_proj = FFNN(input_dim=1280, hidden_dim=4096, output_dim=2048)
traj_proj  = FFNN(input_dim=1792, hidden_dim=4096, output_dim=2048)
loss_fn    = InfoNCELoss(temperature_init=0.07)

v = torch.randn(32, 1280)  # batch of voice embeddings
t = torch.randn(32, 1792)  # paired trajectory embeddings

z_v = voice_proj(v)
z_t = traj_proj(t)
loss = loss_fn(z_v, z_t)

Quick example: bijective round-trip

import torch
from joint_embedding_voice2traj import RealNVP1D

flow = RealNVP1D(dim=2048, n_blocks=8, hidden_dims=[256, 256])
x = torch.randn(16, 2048)

z, log_det = flow.forward(x)
x_rec, _   = flow.inverse(z)

assert torch.allclose(x, x_rec, atol=1e-5)

Quick example: ATC transcription

from joint_embedding_voice2traj import ATCTranscriptor

transcriptor = ATCTranscriptor(use_lm=False, output_hidden_states=False)
result = transcriptor.forward("path/to/atc_audio.mp3", return_representation=True)

print(result["text"])
print(result["representation"].shape)  # (1280,) with TIME_MEAN_POOL=True

---

Testing

Tests are organized under src/tests/ by subsystem.

Windows PowerShell:

.\dev\test.ps1 run-all
.\dev\test.ps1 run-real-nvp
.\dev\test.ps1 run-x-lstm
.\dev\test.ps1 run-lstm-normal
.\dev\test.ps1 run-ffnn-contrastive
.\dev\test.ps1 run-wav2vect
.\dev\test.ps1 clean-all

Unix / macOS:

make test-run-all
make test-run-real-nvp
make test-run-x-lstm
make test-run-lstm-normal
make test-run-ffnn-contrastive
make test-run-wav2vect
make test-clean-all

Direct pytest:

uv run pytest -v src/tests/
uv run pytest -v src/tests/test_real_nvp/

Notes:

• test_wav2vect uses a real MP3 file and a pretrained model — it is heavier than the rest.
• Coverage HTML artifacts are written to htmlcov/.
• See src/tests/TESTING.md for detailed guidance per suite.

---

PDF Reports

Generated LaTeX reports in pdf_latex_reports/:

---

Sources and References

Research context

• CITRIS Aviation: https://citris-uc.org/

xLSTM

• Beck, M. et al. (2024). xLSTM: Extended Long Short-Term Memory. arXiv: https://arxiv.org/abs/2405.04517

RealNVP and normalizing flows

• Dinh, L., Sohl-Dickstein, J., Bengio, S. (2017). Density Estimation using Real-NVP. arXiv: https://arxiv.org/abs/1605.08803

Contrastive learning

• van den Oord, A., Li, Y., Vinyals, O. (2018). Representation Learning with Contrastive Predictive Coding. arXiv: https://arxiv.org/abs/1807.03748

ATC speech model

• Zuluaga-Gomez, J. et al. (2022). How to Improve Automatic Speech Recognition for Air Traffic Control. arXiv: https://arxiv.org/abs/2203.16822
• HuggingFace model: https://huggingface.co/Jzuluaga/wav2vec2-xls-r-300m-en-atc-uwb-atcc-and-atcosim
• Wav2Vec2 / XLS-R: https://huggingface.co/facebook/wav2vec2-xls-r-300m

Tooling

• PyTorch: https://pytorch.org/
• Transformers: https://huggingface.co/docs/transformers/index
• Torchaudio: https://pytorch.org/audio/stable/index.html
• faster-whisper: https://github.com/SYSTRAN/faster-whisper
• UMAP: https://umap-learn.readthedocs.io/