Trajectory-balance normalization experiments

[josephdviviano]

Summary

• Implements and benchmarks trajectory-length normalization for GFlowNet TB and VarGrad losses on HyperGrid environments, testing whether dividing squared scores by trajectory length T improves convergence when trajectory lengths vary.
• Adds a full experiment pipeline: 4 custom HyperGrid reward environments (original, cosine, bitwise_xor, multiplicative_coprime), Optuna HP search (50 trials per algo×env), and confirmation runs (top-3 configs × 5 seeds × 4000 iterations).
• Restructures the repo: renames code/ to lir/, adds conftest.py, environment.yaml, and a test suite for LIR gradient correctness, convergence, and pruning.

What changed

• lir/gflownet/: tb_normalize.py (training loop, 4 algorithms, CLI), hypergrid.py (ModifiedHyperGrid with 4 reward functions, mode analysis, GF(2) feasibility checks), checkpoint.py (atomic run bookkeeping).
• experiments/: optuna_sweep.py (two-phase Optuna driver), SLURM launch scripts, gflownet_results.ipynb (all figures and tables for the paper).
• lir/test/: Gradient sign verification, fixed-point convergence, attention-mask pruning tests.
• lir/lir__simpler.py: Core LIR training loop (moved from code/, expanded).
• WRITEUP.md: Synthesized experiment notes with key findings on gradient clipping and logZ learning rate behavior.

Key results

• ModLPV removes the need for extreme gradient clipping: LPV requires grad_clip ~0.01–0.04 (search boundary), while ModLPV works at ~0.2–1.2 (10–65× less aggressive).
• ModTB requires faster logZ learning on complex environments: effective logZ lr is ~46× higher for ModTB vs TB on multiplicative_coprime (20k modes), but similar on original (256 modes).

Test plan

• pytest passes (core LIR tests + checkpoint tests)
• python -m lir.gflownet.tb_normalize --envs original --algos TBGFlowNet --n_iterations 10 --n-seeds 1 runs without error
• Notebook experiments/gflownet_results.ipynb renders against existing optuna_results/

🤖 Generated with Claude Code