VAE-LSTM: Pose-State Discovery Pipeline

This repo trains a fused VAE -> LSTM -> GMM model for unsupervised behavioral state discovery from pose keypoints.

What it does

• Preprocesses pose detections from a structured CSV (keypoint x/y/conf for 6 keypoints, optional ts_ms).
• Trains a shallow VAE on single frames to learn a compact latent representation.
• Trains an LSTM over latent sequences and fits a GMM on LSTM hidden states for state labels.
• Exports a fused model to ONNX for TensorRT deployment.

Data format

• Input CSV: data/detections.csv (example in repo).
• Expected columns for structured input:
  • ts_ms (optional but required for LSTM gap-aware windowing)
  • nose_x, nose_y, nose_conf
  • head_x, head_y, head_conf
  • left_ear_x, left_ear_y, left_ear_conf
  • right_ear_x, right_ear_y, right_ear_conf
  • back_x, back_y, back_conf
  • tailbase_x, tailbase_y, tailbase_conf
• Flat x,y,conf triplet CSVs are also supported via a fallback loader.

Preprocessing (actual pipeline)

• Confidence gating: low-confidence keypoints are held from the previous valid frame; missing flags are appended.
• Canonicalization: per-frame translation and scale normalization.
  • Center: midpoint of head/back when available, else mean of available keypoints.
  • Scale: distance head<->back, else tailbase<->back, else sqrt(bbox area), else 1.0.
  • No rotation alignment is performed.
• Features:
  • Per keypoint x,y (12 dims)
  • Angle features (sin/cos) for all pairwise keypoint combinations (15 pairs, 30 dims)
  • Missing flags (6) + rot_quality (1)
  • Total dims 49; VAE input excludes the last 7 flag/rot dims.
• Normalization: z-score continuous features with clip to +/-5 std; flags/rot are not normalized.

Training

• Stage 1 VAE (single CSV):
  • python train_vae.py --csv data/detections.csv
  • Saves artifacts/vae.pt, artifacts/vae_best.pt, and artifacts/stats.npz.
• Stage 1 VAE (multi-session):
  • python train_vae.py --data-root data
  • Discovers data/*/detections.csv and computes global normalization stats across sessions.
• Stage 2 LSTM + GMM (single CSV):
  • python train_lstm.py --csv data/detections.csv
  • Requires timestamps in the CSV (ts_ms) or via --timestamps-csv.
  • Saves artifacts/fused.pt and GMM params (gmm_mu.npy, gmm_precision.npy).
• Stage 2 LSTM + GMM (multi-session):
  • python train_lstm.py --data-root data --shuffle-sessions --session-seed 0
  • Splits by session (default ratios 0.7/0.15/0.15) and builds gap-aware windows per session.

Export

• python export_onnx.py --fused-ckpt artifacts/fused.pt --gmm-mu artifacts/gmm_mu.npy --gmm-precision artifacts/gmm_precision.npy
• Output: artifacts/model.onnx

Utilities

• tools/extract_state_clips.py slices top bouts per state into video clips using validation labels and window starts.

Notes

• The pipeline does not compute velocity/acceleration features.
• LSTM training requires gap-aware windows, so timestamps are required.