LoopedHRM

A research codebase for Two-Level Adaptive Compute — Hierarchical Iterative Solver with Sudoku as the first environment.

This repository is intentionally scaffolded as a skeleton. It captures project structure, conceptual interfaces, and implementation placeholders while preserving the core design goals:

• iterative latent refinement (Looped-style)
• hierarchical reasoning stages (HRM-style)
• learned halting at both refinement and segment levels
• adaptive compute allocation based on puzzle difficulty

Project Status

🚧 Early scaffold stage. Modules currently contain interfaces and TODO markers.

Core Research Goal

Learn how much computation to use while solving structured reasoning tasks.

Two compute levels are modeled as interacting modules:

1. Lower module for local latent refinement (variable refine-step bursts)
2. Upper module for global update, evaluation, and continuation/stop decisions

The model should learn to stop both levels dynamically and allocate lower/upper compute based on puzzle difficulty.

Repository Layout

.
├── configs/                     # Experiment and model configuration skeletons
├── docs/
│   └── design_doc.md            # Living design document from the proposal
├── notebooks/                   # Analysis notebooks (future)
├── scripts/                     # Training/evaluation entrypoints
├── src/looped_hrm/
│   ├── data/                    # Sudoku generation/loading and batching
│   ├── models/                  # Hierarchical iterative solver components
│   ├── compute/                 # Halting and compute budget policies
│   ├── training/                # Losses, trainer loops, curriculum
│   ├── evaluation/              # Metrics and benchmarking
│   └── utils/                   # Common helpers
└── tests/                       # Unit/integration test skeletons

Quick Start (uv)

uv sync --extra dev
uv run python scripts/train.py --config configs/experiment/sudoku_baseline.yaml

The training pipeline is not yet implemented; current scripts are placeholders.
uv run ... automatically uses the project environment, so manual activation is optional.

Sudoku Dataset Source

This project now expects Sudoku data in the HRM dataset format (generated by dataset/build_sudoku_dataset.py from https://github.com/sapientinc/HRM).

Place the generated data folder under data/ (for example data/sudoku-extreme-full) so looped_hrm.data.sudoku_dataset.SudokuDataset can load train/ and test/ splits. Alternatively all scripts that use a dataset can be pointed to a directory.

Pre-merge Testing Workflow

Use this workflow before merging dataset-loader changes.

1. Install dependencies (runtime + dev/test):

   uv sync --extra dev

2. Run unit tests:

   uv run pytest -q

3. Run synthetic smoke test (no external data required):

   uv run python scripts/validate_sudoku_dataset.py --self-test

4. Validate against a real HRM dataset folder (if available):

   uv run python scripts/validate_sudoku_dataset.py --dataset-path data/sudoku-extreme-full

## Development Priorities

1. Implement Sudoku dataset and difficulty buckets.
2. Implement latent state update/refinement module.
3. Implement dual halting heads:
   - refinement halt
   - segment halt
4. Add training objective balancing correctness vs compute penalty.
5. Instrument compute-usage metrics by difficulty.

## Design Source

See `docs/design_doc.md` for the full conceptual plan and constraints.