Skip to content

Folders and files

NameName
Last commit message
Last commit date

Latest commit

 

History

52 Commits
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 

Repository files navigation

GENESIS

Learning to select search algorithms under a joint compute + physical budget.

GENESIS studies a resource-rational question: when a controller is given a portfolio of complete search algorithms (A*, Dijkstra, greedy best-first) and a finite energy budget, does reinforcement learning teach it to pick the right algorithm for each map — cheap search when a rough answer suffices, expensive optimal search only when it pays off? It's the orchestration problem behind a routing engine ("which planner do I run for this query?"), studied in a fast, fully-controllable 2D grid world.

A PPO controller (it orchestrates search; it does not navigate) chooses among:

action effect cost
Plan-A* / Plan-Dijkstra / Plan-Greedy run that algorithm over the known map (budget-capped) → store a route energy = nodes expanded × PLAN_UNIT
Follow walk the stored route on the true grid; re-plan on a wall/terrain surprise energy per step
Sense reveal a disc of the map (under fog) flat
Idle wait small "life-support" energy

Maps are 128×128, procedurally generated (random / rooms / maze / trap "comb" families), with weighted terrain (rough patches cost more to walk) and an optional fog of war. Because the energy budget can't afford a full optimal search on every map, the algorithm choice matters — that is the whole experiment.

This is v2. The original design (13 sub-algorithmic atoms + a PAIRED adversarial curriculum) is preserved in the codebase and history but was superseded — see CLAUDE.md for the full arc and the reasons for the redesign.

Key finding (honest, in progress)

The central result is mechanistic and conditional, and it's the interesting part:

  • Under fog, algorithm selection is impossible — and the agent is right not to try. The planner sees terrain only where it has sensed; everywhere else it assumes flat, cost-1 ground. So a plan over a mostly-unknown map is a plan over a flat world, where A* and greedy return nearly the same route and A* merely pays 2–3× the expansions. A 6000-iteration agent reached 76% success but used Plan-A* just 0.3% of the time, with zero correlation between map roughness and A* usage. It correctly converged to the cheapest planner.
  • With full visibility, selection becomes learnable. Measured at 128×128 (total energy = plan + walk): on clean maps A* wins (116 vs greedy's 179 — it finds real routes around walls); on very rough maps optimal search blows the compute budget and greedy wins. A warm-started fog-off agent immediately shifts to using A* more than greedy — the first genuine selection signal.

So the emerging story is "learned algorithm selection emerges only when the cost structure is observable, and provably cannot under partial information" — a positive result and a clean negative, with the environment ablations (fog on/off, terrain on/off) already run as controls. The decisive metric — the per-map roughness↔A*-usage correlation on the finished fog-off run — is being measured; see docs/RESUME.md.

Stack

JAX (jax[cuda12]), Equinox, Optax, Distrax, chex — no PyTorch/TF in the core loop. Everything is JAX-native (jit + vmap, fixed-shape lax.scan), so the environment and the search algorithms batch across hundreds of parallel worlds on one GPU. Native Python 3.12 in a WSL2 (Ubuntu-24.04) venv.

Quick start

source .venv/bin/activate                       # deps pinned in requirements.lock.txt
JAX_PLATFORMS=cpu python -m pytest -q            # 125 tests (env, search kernel, portfolio, training)

# train the v2 portfolio controller at 128x128 (fog off = the selection-learnable regime)
GENESIS_SIZE=128 GENESIS_ENVS=32 GENESIS_UPD=2000 GENESIS_FOG=0 \
XLA_PYTHON_CLIENT_PREALLOCATE=false python -m genesis.train_v2

# watch it live in a browser (episode replays + per-algorithm usage; nothing saved to disk)
python -m genesis.watch runs/v2_128_nofog/ckpt --fog       # -> http://localhost:8000

Run knobs are environment variables: GENESIS_SIZE, GENESIS_ENVS, GENESIS_UPD, GENESIS_FOG (0 = off), GENESIS_WARM=<ckpt> (carry a trained agent into a changed environment), GENESIS_TAG. Training checkpoints every eval and resumes the same command after a crash; changing the environment config auto-starts fresh (a fingerprint guard).

Repository layout

genesis/
  config.py         EnvConfig — the static world definition (size, budget, terrain, fog, topologies)
  generate.py       solvable-by-construction map generation (BFS flood-fill) + topology families
  portfolio.py      the generic budget-capped best-first kernel (A*/Dijkstra/greedy, one scan)
  portfolio_env.py  v2 environment: PortfolioState + Plan/Follow/Sense/Idle action space
  train_v2.py       PPO training driver (resume, warm-start, metrics logging)
  watch.py          live browser dashboard (episode replays + algorithm-usage analytics)
  network/ppo/…     Equinox actor-critic, clipped PPO, GAE, running obs-normalizer, instrumentation
  primitives.py, env.py, paired.py, …   the v1 system (13 atoms + PAIRED), kept for reference/ablation
tests/              125 tests — solvability, search optimality vs BFS, budget accounting, training smokes
docs/               RESUME.md (handoff), the design proposals, and the research roadmap

Where to go next

  • docs/RESUME.md — one-page state + prioritized next actions (baselines & VBS, multi-seed, the correlation measurement). Start here to continue the work.
  • CLAUDE.md — the full design brief, locked decisions, and the complete development history.
  • docs/Portfolio_Redesign_Proposal_2026-07.md — the v2 design rationale and how "beats all algorithms" is measured (Virtual Best Solver).

Status

Environment, all v2 primitives, the search portfolio, PPO training (with resume/warm-start/live dashboard), and the v1 system are implemented and tested (125 tests green). The open scientific work is the measurement layer — per-algorithm baselines, the Virtual Best Solver oracle, and multi-seed headline runs — and confirming the fog-off selection result. See docs/RESUME.md.

About

Research Project

Resources

Stars

0 stars

Watchers

0 watching

Forks

Releases

Packages

Contributors

Languages