{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "9ed15c04",
   "metadata": {},
   "source": [
    "# Recursive Reasoning Failures are Chaotic\n",
    "\n",
    "Small recursive reasoners (HRM, TRM) iterate a latent state to solve a puzzle (Sudoku) before\n",
    "emitting an answer. **The core finding:** measured along the recurrent inference trajectory, the\n",
    "**leading finite-time Lyapunov exponent (λ₁) is higher on failed examples than on successful ones**\n",
    "— in the *same* trained network. Failure is locally more chaotic.\n",
    "\n",
    "This notebook, top to bottom:\n",
    "1. **Load** a trained model from HuggingFace (`blackhao0426/recursive-reasoning-chaos`).\n",
    "2. **The result** — compute λ₁ per example (Benettin / JVP along the trajectory) and show the\n",
    "   success-vs-failure separation (histogram + AUC). *This is the headline.*\n",
    "3. **Why** — run the recurrence far past its training budget: failures are a *transient* — they\n",
    "   escape the chaotic set and self-correct given enough compute (TRM), or stay trapped much longer\n",
    "   (HRM). Neither settles to a wrong fixed point.\n",
    "4. **Basin accessibility** — restart from a perturbed initial state: is a failure input-determined\n",
    "   or initial-condition-determined?\n",
    "\n",
    "Companion analysis repo: `github.com/YurenHao0426/recursive-reasoning-dynamics`."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f9d8321d",
   "metadata": {},
   "source": [
    "## 0. Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ea63e86f",
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install -q torch einops pydantic huggingface_hub numpy matplotlib tqdm\n",
    "import numpy as np, matplotlib.pyplot as plt, torch\n",
    "from tqdm.auto import tqdm\n",
    "print(\"torch\", torch.__version__, \"| cuda\", torch.cuda.is_available())\n",
    "if not torch.cuda.is_available():\n",
    "    print(\"\\n⚠️  No GPU detected — the JVP/rollout cells will be SLOW on CPU.\")\n",
    "    print(\"   Colab: Runtime → Change runtime type → Hardware accelerator → GPU (T4), then re-run.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1da5260b",
   "metadata": {},
   "source": [
    "## 1. Load a trained model from HuggingFace\n",
    "\n",
    "Downloads model code + checkpoint + a 2000-puzzle test set from `blackhao0426/recursive-reasoning-chaos`.\n",
    "`MODEL` ∈ {`trm_sudoku`, `hrm_sudoku`}. **TRM is MLP-only → runs on a laptop CPU.** To switch\n",
    "models, change `MODEL` and **restart the kernel** (the two ship same-named `models` packages)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d1821039",
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys, yaml, json\n",
    "from pathlib import Path\n",
    "from huggingface_hub import snapshot_download\n",
    "\n",
    "HF_REPO = \"blackhao0426/recursive-reasoning-chaos\"\n",
    "MODEL = \"trm_sudoku\"   # or \"hrm_sudoku\"\n",
    "root = Path(snapshot_download(HF_REPO))\n",
    "sys.path.insert(0, str(root / (\"code_trm\" if MODEL.startswith(\"trm\") else \"code_hrm\")))\n",
    "\n",
    "cfg = yaml.safe_load((root / MODEL / \"all_config.yaml\").read_text())\n",
    "meta = json.loads((root / \"data\" / \"sudoku_meta.json\").read_text())\n",
    "arch = dict(cfg[\"arch\"]); arch.update(batch_size=64, seq_len=meta[\"seq_len\"], vocab_size=meta[\"vocab_size\"],\n",
    "                                       num_puzzle_identifiers=meta[\"num_puzzle_identifiers\"], causal=False)\n",
    "if MODEL.startswith(\"trm\"):\n",
    "    from models.recursive_reasoning.trm import TinyRecursiveReasoningModel_ACTV1 as M\n",
    "else:\n",
    "    from models.hrm.hrm_act_v1 import HierarchicalReasoningModel_ACTV1 as M\n",
    "model = M(arch)\n",
    "sd = torch.load(root / MODEL / \"weights.pt\", map_location=\"cpu\", weights_only=True)\n",
    "model.load_state_dict({k.replace(\"_orig_mod.\",\"\").replace(\"model.\",\"\"): v for k,v in sd.items()}, strict=False)\n",
    "dev = \"cuda\" if torch.cuda.is_available() else \"cpu\"; model.to(dev).eval()\n",
    "inner = model.inner\n",
    "inp = np.load(root/\"data\"/\"sudoku_test_inputs.npy\"); lab = np.load(root/\"data\"/\"sudoku_test_labels.npy\")\n",
    "pid = np.load(root/\"data\"/\"sudoku_test_pid.npy\")\n",
    "print(f\"loaded {MODEL}: hidden={inner.config.hidden_size}, H_cycles={inner.config.H_cycles}, L_cycles={inner.config.L_cycles}, test puzzles={len(inp)}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e0b10696",
   "metadata": {},
   "source": [
    "## 2. The core result — failures are more chaotic (leading FTLE / λ₁)\n",
    "\n",
    "For each puzzle we run the recurrence for the 16-segment inference budget and propagate one tangent\n",
    "vector through every module update (forward-mode JVP), renormalizing each step and accumulating the\n",
    "log-growth (Benettin's method for the largest exponent). λ₁ = mean log-growth per module-evaluation.\n",
    "Then split by outcome at segment 16. **Failures sit at higher λ₁** — they are locally expanding /\n",
    "chaotic; successes have collapsed toward the solution. `AUC(−λ₁ → success)` near 1 = clean separation.\n",
    "\n",
    "The JVP uses the fact that the update is `module(a, b) = layers(a + b)`, so a perturbation of the\n",
    "combined input can be fed through one slot. (GPU: ~1 min. Laptop/CPU with TRM: a few min — lower `n`.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "01a845b8",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch.autograd.functional as AF\n",
    "from contextlib import nullcontext\n",
    "try:                                   # HRM attention JVP needs the math SDP backend (no FlashAttn double-backward)\n",
    "    from torch.nn.attention import sdpa_kernel, SDPBackend\n",
    "    MATHCTX = lambda: sdpa_kernel(SDPBackend.MATH)\n",
    "except Exception:\n",
    "    MATHCTX = nullcontext\n",
    "\n",
    "def auc(score, y):\n",
    "    p, n = score[y==1], score[y==0]\n",
    "    if len(p)==0 or len(n)==0: return float('nan')\n",
    "    a=np.concatenate([p,n]); o=np.argsort(a); r=np.empty(len(a)); r[o]=np.arange(1,len(a)+1)\n",
    "    return (r[:len(p)].sum()-len(p)*(len(p)+1)/2)/(len(p)*len(n))\n",
    "\n",
    "def leading_ftle(inp, lab, pid, n=128, n_seg=16, seed=0):\n",
    "    rng=np.random.default_rng(seed); idx=rng.choice(len(inp), n, replace=False)\n",
    "    pe=inner.puzzle_emb_len; sf=inner.config.seq_len+pe; hid=inner.config.hidden_size; D=sf*hid; B=n\n",
    "    is_hrm=hasattr(inner,\"H_level\") and getattr(inner,\"H_level\",None) is not None\n",
    "    Hmod=inner.H_level if is_hrm else inner.L_level                 # weight-tied TRM reuses L_level\n",
    "    X=torch.tensor(inp[idx].astype(np.int32),device=dev); Y=torch.tensor(lab[idx].astype(np.int32),device=dev)\n",
    "    P=torch.tensor(pid[idx].astype(np.int32),device=dev)\n",
    "    si=dict(cos_sin=inner.rotary_emb() if hasattr(inner,\"rotary_emb\") else None)\n",
    "    g=torch.Generator(device=dev).manual_seed(seed)\n",
    "    jvp=lambda f,x,v: AF.jvp(f, x, v=v, create_graph=False, strict=False)\n",
    "    def renorm(vH,vL):\n",
    "        nrm=torch.sqrt(vH.pow(2).sum(1,keepdim=True)+vL.pow(2).sum(1,keepdim=True)).clamp_min(1e-30)\n",
    "        return vH/nrm, vL/nrm, nrm.squeeze(1)\n",
    "    with MATHCTX():\n",
    "        emb=inner._input_embeddings(X,P); m=Y>0\n",
    "        zH=inner.H_init.unsqueeze(0).expand(B,sf,hid).clone().to(inner.forward_dtype)\n",
    "        zL=inner.L_init.unsqueeze(0).expand(B,sf,hid).clone().to(inner.forward_dtype)\n",
    "        vH=torch.randn(B,D,device=dev,generator=g); vL=torch.randn(B,D,device=dev,generator=g)\n",
    "        vH,vL,_=renorm(vH,vL); logsum=torch.zeros(B,device=dev); nstep=0\n",
    "        for seg in tqdm(range(n_seg), desc=\"FTLE (segments)\"):\n",
    "            with torch.enable_grad():\n",
    "                zH,zL=zH.detach(),zL.detach()\n",
    "                for _h in range(inner.config.H_cycles):\n",
    "                    for _l in range(inner.config.L_cycles):\n",
    "                        vc=(vH+vL).reshape(B,sf,hid).to(inner.forward_dtype)\n",
    "                        zL,Dv=jvp(lambda z: inner.L_level(z, zH+emb, **si), zL, vc)\n",
    "                        vL=Dv.reshape(B,D).float(); vH,vL,grow=renorm(vH,vL); logsum+=grow.log(); nstep+=1\n",
    "                    vc=(vH+vL).reshape(B,sf,hid).to(inner.forward_dtype)\n",
    "                    zH,Dv=jvp(lambda z: Hmod(z, zL, **si), zH, vc)\n",
    "                    vH=Dv.reshape(B,D).float(); vH,vL,grow=renorm(vH,vL); logsum+=grow.log(); nstep+=1\n",
    "        ftle=(logsum/nstep).cpu().numpy()\n",
    "        ok=(((inner.lm_head(zH)[:,pe:].float().argmax(-1)==Y)|~m).all(-1)).cpu().numpy()\n",
    "    return ftle, ok\n",
    "\n",
    "ftle, succ = leading_ftle(inp, lab, pid, n=64)   # n=64 already separates cleanly; raise for tighter histograms\n",
    "print(f\"success rate {succ.mean():.2f} | median λ1  success {np.median(ftle[succ]):+.4f}  vs  failure {np.median(ftle[~succ]):+.4f}\")\n",
    "print(f\"AUC(-λ1 -> success) = {auc(-ftle, succ.astype(int)):.3f}   (>0.5 means failures are more chaotic)\")\n",
    "plt.figure(figsize=(6,4))\n",
    "b=np.linspace(ftle.min(), ftle.max(), 40)\n",
    "plt.hist(ftle[succ], b, alpha=.6, color='g', density=True, label=f'success (n={succ.sum()})')\n",
    "plt.hist(ftle[~succ], b, alpha=.6, color='r', density=True, label=f'failure (n={(~succ).sum()})')\n",
    "plt.axvline(0, ls=':', c='k', lw=1)\n",
    "plt.xlabel('leading finite-time Lyapunov exponent  λ1'); plt.ylabel('density')\n",
    "plt.title(f'{MODEL}: failures are more chaotic'); plt.legend(); plt.tight_layout(); plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "36e16b8f",
   "metadata": {},
   "source": [
    "## 3. Why — transient chaos: failures *escape* with more compute\n",
    "\n",
    "Run the recurrence `N_SEG` segments (far past the 16-segment budget) and watch the fate of\n",
    "trajectories that fail at segment 16. **TRM** failures escape the chaotic transient and resolve to\n",
    "the correct answer; **HRM** failures are far more strongly trapped. Re-run cell 1 with\n",
    "`MODEL=\"hrm_sudoku\"` (restart kernel) to compare."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6603b85d",
   "metadata": {},
   "outputs": [],
   "source": [
    "def extended_rollout(inp, lab, pid, n=256, n_seg=128, seed=0):\n",
    "    rng=np.random.default_rng(seed); idx=rng.choice(len(inp), n, replace=False)\n",
    "    pe=inner.puzzle_emb_len; sf=inner.config.seq_len+pe; hid=inner.config.hidden_size\n",
    "    is_hrm=hasattr(inner, \"H_level\") and getattr(inner,\"H_level\",None) is not None\n",
    "    Hmod=inner.H_level if is_hrm else inner.L_level\n",
    "    X=torch.tensor(inp[idx].astype(np.int32),device=dev); Y=torch.tensor(lab[idx].astype(np.int32),device=dev)\n",
    "    P=torch.tensor(pid[idx].astype(np.int32),device=dev)\n",
    "    EX=[]; DR=[]\n",
    "    with torch.no_grad():\n",
    "        zH=inner.H_init.unsqueeze(0).expand(n,sf,hid).clone().to(inner.forward_dtype)\n",
    "        zL=inner.L_init.unsqueeze(0).expand(n,sf,hid).clone().to(inner.forward_dtype)\n",
    "        si=dict(cos_sin=inner.rotary_emb() if hasattr(inner,\"rotary_emb\") else None)\n",
    "        emb=inner._input_embeddings(X,P); m=Y>0; prev=None\n",
    "        for _ in tqdm(range(n_seg), desc=\"rollout (segments)\"):\n",
    "            for _h in range(inner.config.H_cycles):\n",
    "                for _l in range(inner.config.L_cycles):\n",
    "                    zL=inner.L_level(zL, zH+emb, **si)\n",
    "                zH=Hmod(zH, zL, **si)\n",
    "            p=inner.lm_head(zH)[:,pe:].float().argmax(-1)\n",
    "            EX.append(((p==Y)|~m).all(-1).float().cpu().numpy())\n",
    "            DR.append((torch.zeros(n) if prev is None else (zH-prev).float().flatten(1).norm(dim=1).cpu()).numpy())\n",
    "            prev=zH.detach()\n",
    "    return np.stack(EX,1), np.stack(DR,1)\n",
    "\n",
    "ex, dr = extended_rollout(inp, lab, pid, n=256, n_seg=128)\n",
    "T=ex.shape[1]; fail=ex[:,15]==0; nf=fail.sum()\n",
    "print(f\"accuracy @16={ex[:,15].mean():.3f}  @{T}={ex[:,-1].mean():.3f}\")\n",
    "print(f\"of {nf} step-16 failures: self-resolve to CORRECT by seg{T}: {(fail&(ex[:,-1]==1)).sum()/max(nf,1)*100:.0f}%\")\n",
    "fig,ax=plt.subplots(1,2,figsize=(11,4))\n",
    "ax[0].plot(range(1,T+1), ex.mean(0)); ax[0].axvline(16,ls='--',c='gray'); ax[0].set_xscale('log')\n",
    "ax[0].set_xlabel('inference segments'); ax[0].set_ylabel('accuracy'); ax[0].set_title('accuracy vs compute')\n",
    "S=[(fail&(ex[:,:s].max(1)==0)).sum()/max(nf,1) for s in range(16,T+1)]\n",
    "ax[1].plot(range(16,T+1),S); ax[1].set_yscale('log'); ax[1].set_xlabel('segments'); ax[1].set_ylabel('frac failures still unsolved')\n",
    "ax[1].set_title('escape from chaotic set (straight line on log-y = exponential escape)'); plt.tight_layout(); plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "32c11219",
   "metadata": {},
   "source": [
    "## 4. Basin accessibility — input-determined or initial-condition-determined?\n",
    "\n",
    "The puzzle is re-injected every segment (`z_H + input_embeddings`), so perturbing only the\n",
    "**initial** latent `z0` is a clean initial-condition change that leaves the input intact. Restart\n",
    "each step-16 failure `K` times from `z0 + sigma*noise`: if a small kick frees it, the solution basin\n",
    "is large and accessible (TRM); if no nearby IC escapes, the trapping is input-determined (HRM has a\n",
    "hard core). (GPU: seconds. Laptop/CPU with TRM: a couple of minutes — lower `n`/`K`.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "53a705d1",
   "metadata": {},
   "outputs": [],
   "source": [
    "def perturb_z0(inp, lab, pid, n=96, K=8, sigmas=(0.0, 0.1, 0.3, 1.0), n_seg=48, readout=16, seed=0):\n",
    "    rng=np.random.default_rng(seed); idx=rng.choice(len(inp), n, replace=False)\n",
    "    pe=inner.puzzle_emb_len; sf=inner.config.seq_len+pe; hid=inner.config.hidden_size\n",
    "    is_hrm=hasattr(inner,\"H_level\") and getattr(inner,\"H_level\",None) is not None\n",
    "    Hmod=inner.H_level if is_hrm else inner.L_level\n",
    "    sc=float(inner.H_init.float().std()); g=torch.Generator(device=dev).manual_seed(seed)\n",
    "    X=torch.tensor(inp[idx].astype(np.int32),device=dev); Y=torch.tensor(lab[idx].astype(np.int32),device=dev)\n",
    "    P=torch.tensor(pid[idx].astype(np.int32),device=dev)\n",
    "    si=dict(cos_sin=inner.rotary_emb() if hasattr(inner,\"rotary_emb\") else None)\n",
    "    solve=np.zeros((n,len(sigmas),K),bool); base=None\n",
    "    with torch.no_grad():\n",
    "        emb0=inner._input_embeddings(X,P); m0=Y>0\n",
    "        for sj,sg in enumerate(tqdm(sigmas, desc=\"basin (sigma levels)\")):\n",
    "            emb=emb0.repeat_interleave(K,0); Yr=Y.repeat_interleave(K,0); mr=m0.repeat_interleave(K,0); B=n*K\n",
    "            zH=inner.H_init.unsqueeze(0).expand(B,sf,hid).clone().to(inner.forward_dtype)\n",
    "            zL=inner.L_init.unsqueeze(0).expand(B,sf,hid).clone().to(inner.forward_dtype)\n",
    "            if sg>0:\n",
    "                zH=(zH.float()+sg*sc*torch.randn(zH.shape,generator=g,device=dev)).to(inner.forward_dtype)\n",
    "                zL=(zL.float()+sg*sc*torch.randn(zL.shape,generator=g,device=dev)).to(inner.forward_dtype)\n",
    "            solved=torch.zeros(B,dtype=torch.bool,device=dev)\n",
    "            for s in range(n_seg):\n",
    "                for _h in range(inner.config.H_cycles):\n",
    "                    for _l in range(inner.config.L_cycles): zL=inner.L_level(zL,zH+emb,**si)\n",
    "                    zH=Hmod(zH,zL,**si)\n",
    "                ok=((inner.lm_head(zH)[:,pe:].float().argmax(-1)==Yr)|~mr).all(-1); solved|=ok\n",
    "                if sj==0 and s==readout-1: base=ok.view(n,K)[:,0].cpu().numpy()\n",
    "            solve[:,sj]=solved.view(n,K).cpu().numpy()\n",
    "    return solve, base, np.array(sigmas)\n",
    "\n",
    "solve, base, sg = perturb_z0(inp, lab, pid)\n",
    "fail=~base; nf=int(fail.sum())\n",
    "print(f\"{nf} of {len(base)} puzzles fail@16; freeing them by restarting from a perturbed IC:\")\n",
    "for j,s in enumerate(sg):\n",
    "    sub=solve[fail,j]; print(f\"  sigma={s:.1f}: single-restart={sub.mean():.2f}   best-of-K={sub.any(1).mean():.2f}\")\n",
    "plt.figure(figsize=(6,4))\n",
    "plt.plot(sg,[solve[fail,j].mean() for j in range(len(sg))],'o--',label='single restart')\n",
    "plt.plot(sg,[solve[fail,j].any(1).mean() for j in range(len(sg))],'s-',label='best-of-K')\n",
    "plt.xlabel('relative IC noise sigma'); plt.ylabel('solve rate (failing puzzles)')\n",
    "plt.title('basin accessibility: does a restart free a trapped puzzle?'); plt.legend(); plt.grid(alpha=.3); plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6ca89f3b",
   "metadata": {},
   "source": [
    "## What this shows\n",
    "- **The result (cell 2):** in the same trained network, failed trajectories have a higher leading\n",
    "  finite-time Lyapunov exponent than successful ones — failure is locally more chaotic.\n",
    "- **Why (cell 3):** that chaos is a *transient*. Failures sit on a chaotic **saddle**, not a wrong\n",
    "  fixed point — TRM's escape and self-correct with more compute; HRM's are much more strongly\n",
    "  trapped (still a saddle, just a far smaller escape rate). The per-segment escape gap is mostly\n",
    "  compute-per-segment (TRM evaluates its module 21×/segment vs HRM 6×; per module-eval the gap is\n",
    "  only ~1.6×). The \"spurious fixed point\" reading from 2D PCA is an artifact of projecting\n",
    "  high-dimensional chaotic wandering.\n",
    "- **Basin (cell 4):** a small initial-condition kick frees most of TRM's trapped puzzles\n",
    "  (IC-determined, large basin); a hard core of HRM's escapes no nearby IC (input-determined).\n",
    "\n",
    "Try: change `MODEL` (restart kernel), `n`/`n_seg`, and compare TRM vs HRM at every step."
   ]
  }
 ],
 "metadata": {},
 "nbformat": 4,
 "nbformat_minor": 5
}