well_poincare_rl
Hierarchical, multi-step, hyperbolic (Poincaré ball) predictor for spatiotemporal scientific fields, with PPO fine-tuning, continual learning (Replay + EWC), and explicit data/checkpoint contracts for multi-contributor use.
What this is for: production-minded RL and representation work on real scientific data — not demos that silently fall back to synthetic trajectories. The same codebase also supports hierarchy recovery on real taxonomic trees (PBDB), as a direct test of whether hyperbolic geometry helps where theory predicts it should.
Current status
| Capability | Status |
|---|---|
| Channel-agnostic encoder (variable C) | 🔒 Done — verified C=2, 11, 47 |
| Per-domain normalizer + same-C replay | 🔒 Done — end-to-end multi-C |
Well adapter (channels-last → (T,C,H,W)) |
🔒 Done — matches real the_well interface |
| Live multi-stream Well (HF) | 🔒 setup / 📋 results — see below |
| Real PBDB hierarchy embedding | 🔒 Done — 2141 nodes, Poincaré > Euclidean at dim=8 |
| Multi-seed / multi-dim real PBDB sweep | Open |
| Attention-based channel fusion | Open |
| Soft hierarchical region priors in the ball | Open (gated on multi-domain latent data) |
Key results
Track A — Real multi-stream continual learning (Well)
Sequential training on live Hugging Face streams, no synthetic fallback.
🔒 Checkpoint-verified: provenance REAL_STREAMED for all three domains,
dataset names (gray_scott_reaction_diffusion, active_matter, shear_flow),
per-domain normalizer channel counts (2, 11, 4 — confirmed via each
normalizer's stored tensor shape), and model architecture (channel-agnostic
encoder, matching this session's redesign, not a stale checkpoint).
📋 Reported (retention table below): the loss values themselves.
run_multistream.py's checkpoint format does not currently save loss/retention
history — only model, params, normalizers, datasets — so these numbers
were not independently re-derived from a saved artifact, only structurally
corroborated by everything above.
| After domain | gray_scott (C=2) | active_matter (C=11) | shear_flow (C=4) |
|---|---|---|---|
| Domain 1 | 0.3486 | — | — |
| Domain 2 | 0.3561 | 0.3299 | — |
| Domain 3 | 0.3248 | 0.3484 | 0.3358 |
Replay buffer by channel count after full run: {2: 24, 11: 24, 4: 24}
(🔒 structurally consistent with the checkpoint's saved normalizer shapes).
Artifact: multistream_continual_gray_scott_active_matter_shear_flow.pt
Track B — Real PBDB hierarchy (Poincaré vs Euclidean)
Live taxonomy edges from Dinosauria + Mammalia (16,000 occurrence records).
🔒 Checkpoint-verified, independently recomputed from raw data: edge hash,
config hash, and identity hash were all recomputed from the raw occurrence
JSON using this repo's own build_edge_list/hash_edge_list/hash_config/
combined_identity_hash functions and matched the checkpoint filename and
its stored results_table — including mean_delta_mrr — to full float
precision. This is the strongest-verified result in the project.
| Setting | ΔMRR | Verification |
|---|---|---|
| dim=8, softmax, c=1.0 | 🔒 +0.1151 | recomputed from checkpoint |
| dim=8, softmax, c=2.0 | 🔒 +0.1576 | recomputed from checkpoint |
Individual Poincaré/Euclidean MRR values (e.g. 0.42 vs 0.305) appear in the
original run log and are arithmetically consistent with the deltas above, but
are 📋 reported, not 🔒 checkpoint-verified — the checkpoint's results_table
stores the delta and win count, not the two absolute MRR values separately.
- Provenance:
REAL_PBDB_TAXONOMY - Scale: 2141 nodes, 2172 edges
- Protocol: reconstruction (train = test edges) — standard for measuring embedding capacity, not link-prediction generalization
- Artifact: content-addressed checkpoint under
checkpoints/+.meta.json
Scope (honest): one seed, one dimension, metrics-only checkpoint (full embedding weights not yet exported). Multi-seed / multi-dim confirmation on this graph is the natural next step, not a re-verification of what's already solid.
Core contracts
Data loading
- Real paths hard-fail on missing data, schema mismatch, or stream failure. They never silently return synthetic data.
- Synthetic data is available only via explicit APIs (
get_synthetic_dataset,--synthetic). Provenance is always reported. - Tensor layout is validated against declared contracts, never guessed from shape heuristics.
Trajectory validation
Checked before training. Trajectories shorter than window + pred_steps
raise TrajectoryTooShortError / EmptyDatasetError instead of being skipped
inside the loop (which previously could train an epoch on zero batches).
Dataset-reuse registry (DatasetRegistry)
Prevents retraining on an already-consumed dataset under multi-contributor use.
(absent) --claim()--> IN_PROGRESS --mark_consumed()--> CONSUMED
\--mark_failed()--> FAILED --allow_retry()--> (absent)
- Claims are atomic (
O_CREAT|O_EXCL). CONSUMEDblocks future claims by default (DatasetAlreadyUsedError).- Explicit override:
--allow-dataset-reuse.
Checkpoints (CheckpointStore)
Content-addressed: filename = sha256(config_hash + code_hash + dataset_hash).
- Atomic writes (temp →
os.replace()only after success). - Partial writes (
.ptwithout.meta.json) raiseCheckpointIntegrityError. - Identical
(config, code, dataset)→DUPLICATE_EXISTS, no redundant write.
Structured outcomes
Contract failures raise typed exceptions with stable outcome_code values
(e.g. NO_DATA_DIRECTORY, CHANNEL_COUNT_MISMATCH, ALREADY_CONSUMED).
Architecture (as shipped)
| Component | Role |
|---|---|
MultiScaleEncoder |
Channel-agnostic: shared 1×1 stem per channel → mean fusion → spatial path |
HierarchicalHyperbolicPredictor |
Multi-step prediction in the Poincaré ball |
FieldNormalizer |
Per-domain z-score; hard-rejects channel-count mismatch |
ReplayBuffer |
Same-C (and optional spatial) filtering; clear error if mixed shapes stacked |
WellStreamAdapter |
Converts real the_well samples (channels-last input/output) → {"fields": (T,C,H,W)} |
DiagonalEWC + teacher distillation |
Continual learning across domains |
HierarchyEmbedding |
Poincaré / Euclidean node embeddings for taxonomy trees |
| Physics losses | Spatial/temporal consistency + generalized channel-coupling (all pairs) |
Design note on variable channels: ingestion (encoder + normalizer + replay) is solved separately from latent organization. Soft hierarchical region priors in the ball remain open and are gated on multi-domain data actually reaching a shared latent space.
Running
pip install torch geoopt optuna gymnasium h5py tqdm the_well
# Contract tests
python -m pytest tests/ -v
# Synthetic full pipeline (no network)
python -m src.run_full --synthetic
# Real local Well-format HDF5 (hard-fails if missing/invalid)
python -m src.run_full --data-root ./data/real
# Live multi-stream continual run (requires HF access, e.g. Kaggle)
python -m src.run_multistream \
--datasets gray_scott_reaction_diffusion active_matter shear_flow \
--max-samples 96 --epochs-per-domain 3
# Hierarchy embedding — synthetic (no network)
python -m src.run_hierarchy_embed --synthetic --tree-type balanced \
--branching-factor 3 --tree-depth 6 --dims 2 3 5 8 \
--loss-types softmax --burn-in-epochs 0 --c-values 1.0 2.0 \
--epochs 80 --lr 0.02 --seeds 0 1 2
# Hierarchy embedding — real PBDB (requires network)
python -m src.data_pbdb --discover --base-name Dinosauria --limit 5 # first
python -m src.run_hierarchy_embed \
--pbdb-taxa Dinosauria Mammalia \
--dims 8 --loss-types softmax --burn-in-epochs 0 \
--c-values 1.0 2.0 --epochs 80 --lr 0.02 --seeds 0 --optimizer radam
Defaults that matter for Track B with current stability guards:
--burn-in-epochs 0 (ablation showed burn-in reduces ΔMRR under guarded
clip_to_ball / clamp_curvature setup).
Artifacts
| File | What it is |
|---|---|
multistream_continual_gray_scott_active_matter_shear_flow.pt |
Track A model + params + per-domain normalizers + dataset list (no loss history) |
checkpoints/<identity_hash>.pt + .meta.json |
Track B metrics table + provenance (content-addressed) |
Large *_8000.json under data cache |
Raw PBDB fetch cache only — optional, regenerable |
Hierarchy checkpoints currently store results_table (metrics), not full
embedding weights. Weight export is a small follow-on if reloadable node
vectors are needed.
Known limitations
- Multi-stream retention numbers are 📋 reported, not 🔒 checkpoint-verified — see Track A above. The run's provenance, channel counts, and architecture are independently confirmed; the specific loss trajectory is not yet, since the current checkpoint format doesn't persist it.
- Real PBDB hierarchy result is one seed / one dimension.
- Channel fusion is mean-pooling (lossy). Attention fusion is the natural upgrade.
- Soft region structure in the Poincaré ball is intentionally deferred until multi-domain data is routinely in a shared latent.
- PPO is implemented and smoke-tested; it is not the primary claim of the current verified results.
- Contributors are assumed to exchange portable checkpoint + meta files; a shared remote registry is a future decision if contributor count grows.
Design principles (non-negotiable)
- Real-data paths hard-fail. Synthetic only via explicit opt-in.
- No silent shape / channel mistakes. Prefer loud, typed errors.
- Provenance is first-class. Every load reports where data came from.
- Checkpoints are content-addressed and atomic.
- Report what the numbers say, including when Euclidean wins or a previously recommended default (e.g. burn-in) is reversed by ablation.
- Verification tier is part of the claim. "Verified" always specifies verified how — checkpoint-recomputed and self-reported are not the same thing, and this document says which is which rather than picking one word for both.
License / packaging
Apache 2.0 intended for model and code release. Package layout is
src/-based and pip-installable for local development.