# hippotorch

[![PyPI](https://img.shields.io/pypi/v/hippotorch?logo=pypi&logoColor=white)](https://pypi.org/project/hippotorch/)
[![Linux CI](https://img.shields.io/github/actions/workflow/status/domezsolt/hippotorch/workflow.yml?branch=main&label=Ubuntu&logo=ubuntu)](https://github.com/domezsolt/hippotorch/actions/workflows/workflow.yml)
[![Windows CI](https://img.shields.io/github/actions/workflow/status/domezsolt/hippotorch/workflow.yml?branch=main&label=Windows&logo=windows)](https://github.com/domezsolt/hippotorch/actions/workflows/workflow.yml)

> **Differentiable episodic memory for reinforcement learning. Retrieves what matters. Forgets what doesn't.**

[Changelog](CHANGELOG.md)

Hippotorch is a drop-in upgrade for replay buffers. It keeps experiences in a learnable memory so agents can remember rare successes, connect distant cause and effect, and transfer knowledge between similar worlds. Under the hood it uses reward-aware contrastive learning, but you mostly interact with a friendly API.

---

## Highlights

- **Memory that adapts with you.** Dual encoders organize episodes by usefulness instead of mere recency.
- **Semantic + uniform sampling.** A single buffer can surface hard-to-find wins while still covering the full state space.
- **Production-friendly extras.** Hugging Face Hub export, FAISS retrieval, Gymnasium wrappers, and health reports ship in the box.
- **Batteries included.** Dozens of scripts and docs show exactly how to benchmark, visualize, and share results.

If you already converge with a plain replay buffer, keep it. Hippotorch shines when agents forget early lessons, face sparse rewards, or operate in partially observed environments.

---

## Installation

```bash
pip install hippotorch            # minimal setup
pip install hippotorch[faiss]     # fast nearest-neighbor retrieval
pip install hippotorch[envs]      # Gymnasium helpers + examples
pip install hippotorch[hub]       # Hugging Face Hub + safetensors
pip install hippotorch[umap]      # projector UMAP export
```

Requirements: Python ≥3.9, PyTorch ≥2.0

---

## Quick Tour

Create an encoder + memory, add episodes, then mix semantic and uniform samples:

```python
import torch
from hippotorch import Episode, DualEncoder, MemoryStore, HippocampalReplayBuffer

state_dim, action_dim = 4, 1
encoder = DualEncoder(input_dim=state_dim + action_dim + 1, embed_dim=128)
memory = MemoryStore(embed_dim=128, capacity=50_000)
buffer = HippocampalReplayBuffer(memory=memory, encoder=encoder, mixture_ratio=0.3)

states = torch.randn(32, state_dim)
actions = torch.randn(32, action_dim)
rewards = torch.randn(32)
buffer.add_episode(Episode(states=states, actions=actions, rewards=rewards))

# Query-aware sampling
query_state = torch.cat([states[0], torch.zeros(action_dim), rewards[:1]])
batch = buffer.sample(batch_size=64, query_state=query_state, top_k=5)

# Sleep/consolidate occasionally
metrics = buffer.consolidate(steps=50, batch_size=64, report_quality=True)
print(metrics["loss"])
```

Rolling with Stable Baselines 3 or Gymnasium? Wrap your existing replay buffer with `SB3ReplayBufferWrapper` or the `HippotorchMemoryWrapper` and keep the rest of your pipeline untouched.

Need hyperparameter guidance? See `docs/diagnostics.md` for health checks and `docs/curriculum.md` for training tips.

---

## Everyday Tools

### Recall While Acting
- Use the lightweight read API: `from hippotorch import query`.
- Pipe `query(..., top_k=5)` results into policies or logging code.
- Gymnasium adapter emits dict observations so SB3 policies can consume retrieval features alongside pixels.
- Examples: `examples/query_inference_demo.py`, `examples/minigrid_memory_wrapper.py`.

### Portable Brains
- Share trained memories with `push_memory_to_hub` / `load_memory_from_hub`.
- Choose local folders for offline passes or Hugging Face Hub for team-wide reuse.
- `scripts/hub_roundtrip_smoke.py` is a 30-second sanity check.
- Docs: `docs/hub.md`.

### Glass-Box Diagnostics
- `buffer.health_report()` returns retrievability, staleness, collapse indicators, and alignment scores.
- Log with `report.to_tensorboard(writer, step)` or `report.to_wandb(run)`.
- See `docs/diagnostics.md` for visuals.

### Batch Retrieval for Low Latency
- `buffer.query_batch(query_vecs, top_k=K)` handles `[B,T,D]` tensors in one go.
- Matches single-query results without looping Python.
- Works with both torch and FAISS backends.

---

## Ready-to-Run Samples

Pick a script, set a seed, and you get a reproducible snapshot:

- **Benchmarks & diagnostics**
  - Retrieval perf: `python scripts/bench_retrieval.py --sizes 10000 100000`
  - Visualization: `python scripts/export_projector_embeddings.py --snapshot run.pt`
  - Retrieval heatmap: `python scripts/retrieval_heatmap.py --memory-checkpoint ...`
- **Environments**
  - CartPole smoke: `bash scripts/quick_cartpole.sh`
  - Corridor curriculum/oracle: `bash scripts/corridor_curriculum.sh`, `bash scripts/corridor_oracle_zn.sh`
  - MiniGrid sweeps: `python scripts/minigrid_memory_benchmark.py --steps 8000 --seeds 3`
  - Intrinsic curiosity example: `python -m examples.intrinsic_demo --episodes 20`
- **Ablations & studies**
  - Rank-weighted consolidation: `bash scripts/run_rank_ablation.sh`
  - Consolidation micro bench: `bash scripts/run_consolidation_micro.sh`
  - Visual MiniGrid clustering: `python -m examples.minigrid_visual --steps 2000`

All scripts keep runtime under a couple of minutes unless stated otherwise. Longer jobs (corridor oracle full run, curriculum sweeps) note their expected duration in the script header.

---

## Learn More

- [docs/benchmarks.md](docs/benchmarks.md) – retrieval setups, FAISS parity, and profiling tips.
- [docs/curriculum.md](docs/curriculum.md) – how to stage corridor tasks and measure regret.
- [docs/usage.md](docs/usage.md) – wrappers, segmenters, and rollout recipes.
- [docs/hub.md](docs/hub.md) – how to move memories between machines or teammates.

Problems or ideas? File an issue, open a discussion, or send a PR.
