Learning to Theorize the World from Observation

Doojin Baek^*, Gyubin Lee^*, Junyeob Baek, Hosung Lee, and Sungjin Ahn

^* Equal contribution

Official PyTorch code for the ICML 2026 Oral paper, recipient of the CompLearn 2026 Best Paper Award.

Neural Theorizer (NEO) instantiates Learning-to-Theorize (L2T) as latent program induction. From raw before-and-after observations, NEO learns reusable primitive operations and composes them into executable theories that transfer to new observations.

Setup

Use Python 3.10 and PyTorch 2.7.1 with CUDA 12.6. The pinned installation is in environments/README.md.

Train and evaluate

The same scripts support gridworld, arithmetic_factorization, and image_editing. For example, from the repository root:

python scripts/generate_data.py --task gridworld --data-root data/gridworld
python scripts/download_checkpoints.py --task gridworld --output-root checkpoints

python scripts/train.py --task gridworld \
  --alpha all --seed all --devices 0 \
  --data-root data/gridworld --output-root runs/gridworld/train \
  --observation-checkpoint checkpoints/gridworld/observation/best_reconstruction/checkpoint_2206.pt

python scripts/evaluate.py --task gridworld \
  --alpha all --seed all --devices 0 \
  --data-root data/gridworld --training-root runs/gridworld/train \
  --output-root runs/gridworld/evaluate

Training runs the three alpha values and seeds 42, 43, 44. Evaluation selects checkpoints by ID query transfer, then evaluates the selected weights on ID, compositional OOD and length OOD. Hard grounding is ON for GridWorld and Arithmetic, and OFF for ImageEditing, in both selection and final evaluation.

Use a new output directory for each run. Add --dry-run to inspect commands. W&B is enabled by default; --wandb-mode offline saves logs locally.

Training and evaluation guide covers all three tasks, observation pretraining, checkpoint paths and NEO-S. ImageEditing requires two training GPUs; GridWorld and Arithmetic require one each. The release provides observation weights; train NEO with the commands above.

Results

Results compares three-seed means and standard deviations with the paper. Per-seed scores are in configs/results.json. This repository implements NEO and NEO-S. Paper baseline values are provided for reference; baseline implementations are not included.

Code

• src/models/: NEO, the theory programmer, executor and vector quantizer.
• src/tasks/: task models, datasets, training and evaluation.
• scripts/: data generation, pretraining, training and evaluation commands.
• configs/: checkpoint manifest and reported results.

Run the tests with CUDA_VISIBLE_DEVICES= python -m pytest -q tests. The CPU test suite also runs on every pull request and push to main. Third-party notices accompany the relevant source.

License

The code and released pretrained observation weights are available under the MIT License. Third-party components retain their original licenses; see third-party notices.

Citation

@inproceedings{baek2026learning,
  title     = {Learning to Theorize the World from Observation},
  author    = {Baek, Doojin and Lee, Gyubin and Baek, Junyeob and Lee, Hosung and Ahn, Sungjin},
  booktitle = {Proceedings of the 43rd International Conference on Machine Learning},
  year      = {2026},
  url       = {https://arxiv.org/abs/2605.03413}
}