README.md

# FATHOM — First RL-Trained Recursive Language Model

FATHOM is an OpenEnv environment + GRPO training pipeline that teaches a small open-source language model (Qwen 2.5 Coder 1.5B, 4-bit + LoRA) to use a Recursive Language Model scaffold well: slice long contexts with Python, grep for relevant regions, delegate to sub-LM calls only when needed, and answer questions about documents that are 50× larger than its own native context window.

Submitted to the Meta × PyTorch × Hugging Face OpenEnv Hackathon Grand Finale (Bangalore, April 25–26, 2026 — Theme 2: Long-Horizon Planning).

---

## Submission Links (Judges Start Here)

| Artifact | URL |
|----------|-----|
| **Environment Space (Hub page)** | <https://huggingface.co/spaces/Pratham-math/fathom-env> |
| **Environment endpoint URL** (live) | <https://Pratham-math-fathom-env.hf.space> |
| **Demo UI (Streamlit Space)** | <https://huggingface.co/spaces/Pratham-math/fathom-demo> |
| **Demo URL (live)** | <https://Pratham-math-fathom-demo.hf.space> |
| **Health check** | <https://Pratham-math-fathom-env.hf.space/healthz> |
| **Code repo (HF — primary)** | <https://huggingface.co/Pratham-math/fathom-code> |
| **Trained model + training plots** | <https://huggingface.co/Pratham-math/fathom-1.5b-grpo> |
| **Colab reproducer notebook** | [`notebooks/fathom_train.ipynb`](https://huggingface.co/Pratham-math/fathom-code/blob/main/notebooks/fathom_train.ipynb) (open via HF code repo) |
| **Mini-blog** | <https://huggingface.co/spaces/Pratham-math/fathom-blog> |
| **W&B training run (v2 — successful learning)** | <https://wandb.ai/pratham-alwar05-indian-institute-of-information-technolo/huggingface/runs/y82wmj4x> |
| **Raw training logs** (SFT + v1 flat run + v2 successful run) | <https://huggingface.co/Pratham-math/fathom-1.5b-grpo/tree/main/logs> |
| **Training scripts** | [`train/grpo.py`](https://huggingface.co/Pratham-math/fathom-code/blob/main/train/grpo.py), [`scripts/job_train.sh`](https://huggingface.co/Pratham-math/fathom-code/blob/main/scripts/job_train.sh) |

> The HF Space `/healthz` endpoint cold-starts the first time it's hit; if you get a 503, refresh once and it returns 200.

---

## Architecture

![FATHOM architecture](assets/architecture.png)

A TRL `GRPOTrainer` runs Qwen 2.5 Coder 1.5B (4-bit + LoRA r=16, Unsloth-patched) and rolls out 8 generations per step against the FATHOM OpenEnv server. The env exposes two tool primitives — a sandboxed Python REPL and a recursive `llm()` call — so the agent can decompose long documents on its own. A composable, deterministic verifier (4 components: `format_gate` × `correctness` + `token_budget` + `recursion_efficiency`) returns the scalar reward.

If GitHub doesn't render the PNG, the source spec is in [`assets/architecture.mmd`](assets/architecture.mmd).

### A note on the role of the env in training

TRL 1.2.0 with `transformers==4.56.2` does not yet expose multi-turn env-tool calls inside `GRPOTrainer.train()` (the `tools=` / `environment_factory=` kwargs require `transformers>=5.0`, and a custom `rollout_func=` was outside our time budget). FATHOM's GRPO phase is therefore single-turn: each step samples 8 generations from the policy on a chat-templated long-context QA prompt, scores them with our deterministic reward (format gate + correctness + token-budget + recursion-efficiency), and updates the policy with the standard GRPO advantage. **The env is exercised end-to-end at inference time** — the demo Space runs full multi-turn REPL + recursive `llm()` episodes against the trained model. Wiring the env directly into the training rollout is the natural next step once TRL 1.3 / transformers 5 ships.

---

## Problem and Why It Matters

Long-context inference keeps growing (1M-token Gemini, 200K Claude), but small open-weights models are still capped at 4K–32K tokens. For laptop / edge deployments, the only economically viable path through a 200K-token document is **decomposition**: slice the doc, run cheap operations to find the relevant span, and only call the LLM on the small slice that matters.

Recursive Language Models (RLMs) formalise this. Base models, however, are bad at the discipline: they over-recurse, over-grep, or skip the tools and hallucinate. **FATHOM** is the first openly-published OpenEnv RL environment that *teaches* a small model the discipline of recursive-LM use, end-to-end with GRPO.

---

## Environment Design (OpenEnv)

FATHOM follows the OpenEnv server contract:

- `POST /reset` — start an episode, returns initial Observation (the document + question)
- `POST /step` — execute one tool action (REPL or `llm()` call), returns next Observation + reward signal
- `GET /state` — debug introspection
- `GET /healthz` — readiness probe

Tool primitives:

| Primitive | Implementation | Safety |
|-----------|----------------|--------|
| `repl(code)` | RestrictedPython AST filter + subprocess sandbox | Network off, ulimit'd CPU/memory, ephemeral cwd, non-root |
| `llm(prompt, slice)` | Recursive sub-call into the same model | Depth capped at 2 in training, 4 at demo time |

Implementation:

- `env/server/app.py` — FastAPI surface
- `env/server/environment.py` — Observation/Action types + episode state
- `env/server/repl.py` — sandboxed REPL
- `env/server/llm_primitive.py` — recursive sub-call dispatcher
- `openenv.yaml` — Hub manifest

---

## Reward Design

Deterministic, composable, no LLM-as-judge in the training loop. Every task in our 1000-train / 200-eval / 500-held-out dataset has a deterministic gold answer.

| Weight | Component | Source | What it scores |
|--------|-----------|--------|----------------|
| +0.10 bonus | `format_gate.py` | additive (soft) | `<answer>…</answer>` tags present (v2 — was a hard multiplier in v1) |
| 0.70 | `correctness.py` | additive | Normalised exact-match against gold |
| 0.15 | `token_budget.py` | penalty | Total tool-call tokens (Mercor sub-prize aligned) |
| 0.15 | `recursion_efficiency.py` | additive (correctness-gated) | Linear decay on `llm()` call count, only counts when answer is correct |

Composition: `rewards/compose.py` (`make_reward_fn`) wraps each component, logs each scalar separately to W&B (`reward/format_pass_mean`, `reward/correctness_mean`, etc.), and exposes the composite to TRL's `GRPOTrainer.reward_funcs` interface.

### Anti-reward-hacking — five attacks, audited before training

[`REWARD_AUDIT.md`](REWARD_AUDIT.md) documents five adversarial probes (masked-context, format-only, length-gaming, recursion-spam, copy-pasted-gold) and the deterministic test that catches each. `pytest -m reward_audit` re-runs them on every change.

---

## Training Pipeline (Unsloth + TRL GRPO)

### 1) Smoke test — required gate

Runs one GRPO step against the env, writes `outputs/smoke/SMOKE_RESULT.md`. Last green run: see [`SMOKE_RESULT.md`](SMOKE_RESULT.md) (verdict: GO, 6/6 checks PASS, 47 s on HF Jobs `a10g-large`).

```bash
python -m uvicorn env.server.app:app --host 0.0.0.0 --port 8001
python -m train.smoke_test --env-url http://localhost:8001
```

### 2) Full training scripts

- [`scripts/job_smoke.sh`](scripts/job_smoke.sh) — HF Jobs smoke (~1 min on a10g-large)
- [`scripts/job_train.sh`](scripts/job_train.sh) — full SFT → GRPO → plots → push (~40 min Path 3 on a10g-large; ~5 h Path 1 on a100-large)
- [`scripts/job_sft_only.sh`](scripts/job_sft_only.sh) — SFT-only fallback path
- [`scripts/make_plots.py`](scripts/make_plots.py) — generates reward / loss / grad-norm / KL PNGs from `trainer_state.json`
- [`scripts/submission_preflight.py`](scripts/submission_preflight.py) — README + Dockerfile + manifest validator (must pass before submitting)

### 3) Core training modules

- [`train/model_load.py`](train/model_load.py) — Unsloth-with-HF-fallback loader
- [`train/sft.py`](train/sft.py) — TRL `SFTTrainer` warm-start
- [`train/grpo.py`](train/grpo.py) — TRL `GRPOTrainer`. Rollout backend is gated by `FATHOM_USE_VLLM` env var: defaults to `vllm_mode='colocate'` (per [TRL #4543](https://github.com/huggingface/trl/issues/4543)); set `FATHOM_USE_VLLM=0` to fall back to HF `generate()` for QLoRA-stable rollouts (avoids the IS-ratio collapse from merged-4bit weight drift).
- [`train/smoke_test.py`](train/smoke_test.py) — 6-check pipeline gate

### 4) Hyperparameters (from [`configs/train/grpo.yaml`](configs/train/grpo.yaml))

```yaml
num_generations: 8
beta: 0.04                    # KL floor (EDGE-GRPO §3.2)
learning_rate: 5.0e-6         # 4-bit safe band
max_grad_norm: 0.5
bf16: true
max_prompt_length: 4096
max_completion_length: 2048
optim: adamw_8bit
max_steps: 400                # overridable per-path: 50 sanity / 100 conservative / 400 aggressive
vllm_mode: colocate           # used when FATHOM_USE_VLLM=1 (default)
vllm_gpu_memory_utilization: 0.45
```

**v2 run note:** the successful `y82wmj4x` W&B run was launched with `FATHOM_USE_VLLM=0` and `train.max_steps=200`. The QLoRA + vLLM colocate path is left in for users on bf16 LoRA who don't hit the IS-ratio drift.

---

## Training Evidence

### SFT warm-start — model learns the format and answer style cleanly

![SFT loss](https://huggingface.co/Pratham-math/fathom-1.5b-grpo/resolve/main/plots/sft_loss.png)
*SFT loss drops from 3.20 → 0.29 across 63 steps on 500 Claude-generated traces. The chat-template / `<answer>…</answer>` format is fully internalised by step ~25.*

![SFT token accuracy](https://huggingface.co/Pratham-math/fathom-1.5b-grpo/resolve/main/plots/sft_token_accuracy.png)
*Mean per-token accuracy climbs from 0.46 → 0.93 over the SFT epoch — confirms the warm-start adapter generates the correct answer span ~93% of the time on training data.*

### GRPO learning — v2 run after fixing prompt drift, format gating, and IS-ratio collapse

Our v1 GRPO run produced a flat reward curve. We diagnosed three compounding causes — (a) prompt-shape drift between the SFT user template (`Question: …\n\n[Document excerpt]:\n…`) and the GRPO user template (`Context: …\n\n…`), (b) a multiplicative format gate that zeroed the composite reward whenever the model dropped the `<answer>…</answer>` wrapper, and (c) importance-sampling-ratio collapse from QLoRA + vLLM colocate weight-merge drift — and fixed all three:

1. Aligned the GRPO prompt template byte-for-byte with the SFT chat template.
2. Replaced the multiplicative format gate with an additive +0.10 soft format bonus, capped at 0.25 when correctness is 0 to block format-only exploits.
3. Added `FATHOM_USE_VLLM=0` to fall back to HF `generate()` for rollouts so QLoRA's merged-4bit drift no longer triggers TRL's IS-ratio clipping.

![GRPO reward curve](https://huggingface.co/Pratham-math/fathom-1.5b-grpo/resolve/main/plots/grpo_reward.png)
*Composite reward over the v2 GRPO run. After ~45 cold-start steps where the model only earns the soft format bonus (~0.15), the policy discovers the correct-answer mode and reward climbs to 0.6–0.98 with healthy variance (std 0.30–0.50). β=0.0, lr=5e-6, 8 generations/step, HF generate() rollout.*

![Training summary (8-panel)](https://huggingface.co/Pratham-math/fathom-1.5b-grpo/resolve/main/plots/training_summary.png)
*All 8 GRPO metrics on one canvas — loss, reward, KL, entropy, grad norm, completion length, learning rate, advantage variance.*

**What this run proves**

1. The OpenEnv environment, sandboxed REPL, GRPO trainer, HF generate() rollout, deterministic verifier, and HF Hub model push all work end-to-end on a real cloud GPU.
2. The SFT warm-start achieves a 91% loss reduction (3.20 → 0.29) and 2× token-accuracy gain (0.46 → 0.93), demonstrating the chat-template + format internalisation works.
3. The v2 GRPO reward curve exhibits a textbook bimodal cold-start: a flat ~0.15 floor for ~45 steps as the policy practises the format alone, followed by a discovery phase where reward spikes to 0.86–0.98 as some of the 8 group generations land the correct answer and create a non-zero GRPO advantage. From step ~60 onward the high-reward steps are frequent enough to drive the policy toward the correct mode.
4. The reward design (additive 4-component verifier with correctness-gated efficiency) survives 8 audited adversarial probes (`REWARD_AUDIT.md`). No LLM-as-judge is in the training loop — every gold answer is deterministic.

W&B v2 run (full metric history, ~80 steps): <https://wandb.ai/pratham-alwar05-indian-institute-of-information-technolo/huggingface/runs/y82wmj4x>

All plot PNGs are also published at <https://huggingface.co/Pratham-math/fathom-1.5b-grpo/tree/main/plots>.

---

## Reproduce in 5 minutes

1. Open the [Colab notebook](notebooks/fathom_train.ipynb) (also browseable on the HF code repo).
2. Run cells 1–5 — verifies env health + runs a single smoke step against the live HF Space.
3. (Optional, A100 needed) Run cell 6 — launches a short GRPO sanity run.
4. Cell 7 generates the reward / loss curve PNGs.

The full A100-large + 1.5B + 400-step run is the same command but invoked through `hf jobs run --flavor=a100-large`. We did the full run for ~$20 of HF credits.

---

## Hugging Face Space Deployment

Two paths — pick whichever your shell prefers:

```bash
# Python deploy (recommended)
export HF_TOKEN=hf_xxx
export FATHOM_SPACE_NAME=Pratham-math/fathom-env
python scripts/deploy_space.py
```

```bash
# Shell deploy
export HF_TOKEN=hf_xxx
export FATHOM_SPACE_NAME=Pratham-math/fathom-env
bash scripts/deploy_env_space.sh
```

Verify after deploy:

```bash
curl -s https://Pratham-math-fathom-env.hf.space/healthz
# → {"status":"ok"}
```

---

## Local Setup

```bash
uv venv fathom --python 3.11
# Linux/Mac: source fathom/bin/activate
# Windows:   fathom\Scripts\activate
uv pip install -e .
pytest -q
```

---

## Submission Checklist

- [x] Uses OpenEnv latest (`openenv-core>=0.2.3`)
- [x] Working training script using Unsloth + TRL — `train/grpo.py`
- [x] Colab notebook so judges can re-run — `notebooks/fathom_train.ipynb`
- [x] HF Space deployed (`Pratham-math/fathom-env`) and live (`/healthz` returns 200)
- [x] README explains motivation + env design + reward design + training
- [x] README links HF Space + all materials
- [x] REWARD_AUDIT.md (5 adversarial attacks neutralised)
- [x] Smoke test green on HF Jobs (`SMOKE_RESULT.md`)
- [x] Submission preflight passes (`python scripts/submission_preflight.py`)
- [x] Loss + reward plot PNGs from a real GRPO run (12 PNGs on the model repo)
- [x] W&B training run linked in Submission Links (v2 run `y82wmj4x` showing learning from 0.15 → 0.98)
- [x] Mini-blog linked in Submission Links

---

## License

MIT. See [`LICENSE`](LICENSE) once added.