Probing Physical Variables in Visual World Models

This repository contains reproducible, notebook-first experiments for asking where physical variables become linearly readable inside visual world models. The current demos cover UMaze, Wall, and PushT and measure position, velocity, acceleration, speed, direction, and object orientation across DINOv2 encoder blocks, learned projected readouts, and dynamics-predictor blocks.

The models are frozen. Each experiment extracts an internal tensor and fits a small ridge-linear map on training trajectories, then evaluates on unseen trajectories and unseen spatial regions. A high score means that the variable is accessible to a linear readout; it does not by itself show that the world model uses that variable.

Start with a notebook

Environment	Demo	Coverage
UMaze	Open Notebook	2,000 unique trajectories
Wall	Open Notebook	1,920 unique trajectories
PushT	Open Notebook	18,500 unique trajectories

*Notebooks will not currently run since they have been anonymized for submission to workshops

The notebooks are the main demos and include explanations, activation extraction, probe fitting, controls, uncertainty estimates, and plots. They clone this repository and download verified assets from releases owned by this repository or from the original public OSF datasets. They do not clone or read from the temporal-straightening repository.

Repository map

notebooks/                 End-to-end Colab demos
scripts/                   Activation extraction and asset utilities
datasets/                  Minimal DINO-WM dataset readers used by the demos
models/                    Minimal DINO-WM runtime needed to load checkpoints
artifacts/checkpoints/     Small configs, checksums, and provenance manifests
results/                   Compact result summaries and selected figures
docs/                      Method and data documentation
tests/                     CPU-only tests for probe construction and assets

Large files are intentionally stored as GitHub Release assets rather than Git objects. Every downloaded checkpoint, dataset, and activation cache is checked against a recorded SHA-256 digest. See Data and artifacts.

Core methodology

For a representation at layer l:

• position uses one frame feature, h[t, l];
• velocity uses the first temporal difference, h[t+1, l] - h[t, l];
• acceleration uses the second temporal difference, h[t+2, l] - 2 h[t+1, l] + h[t, l];
• direction uses unit-vector targets and excludes the slowest training samples;
• PushT orientation uses (cos(theta), sin(theta)) to avoid an angle wrap.

Each representation is standardized using training statistics only. Ridge regression uses lambda=10. Evaluation includes episode holdouts, buffered spatial holdouts, shuffled labels, position-only controls, position-residualized targets, and 1,000 complete-trajectory-window bootstrap resamples for every reported interval. See Methodology.

Local use

git clone --branch initial-release https://github.com/danielhacobian/probing-VLMs.git
cd probing-VLMs
python -m pip install -r requirements.txt

Checkpoint restoration is normally handled by the notebooks. It can also be run directly:

python scripts/fetch_probe_assets.py umaze
python scripts/fetch_probe_assets.py wall

These commands download only from probing-VLMs releases.

Result exports

Each notebook writes validation_selection_scores.csv, headline_selected_test_metrics.csv, headline_straightening_deltas.csv, and headline_protocol.json. The combined exports are committed under results/. After rerunning all three notebooks, refresh them with:

python scripts/collect_headline_exports.py \
  --umaze <UMAZE_OUTPUT_DIR> \
  --wall <WALL_OUTPUT_DIR> \
  --pusht <PUSHT_OUTPUT_DIR>

Reproducibility notes

• Complete trajectories are split 60/20/20 into training, validation, and a locked test partition.
• One deterministic window is selected per trajectory with seed 0.
• OFF and ON conditions use identical windows, labels, splits, and probes.
• Layer/readout selection uses validation trajectories; the selected probe is evaluated once on locked test trajectories.
• Headline values use 1,000 complete-trajectory-window bootstrap resamples and report 95% percentile intervals.
• Notebook outputs currently cover one frozen training seed per condition; bootstrap intervals do not measure training-seed uncertainty.
• UMaze and Wall predictor-action inputs retain the documented legacy-padding limitation; encoder and projected-readout results are unaffected.

Attribution

The minimal model and dataset runtime is adapted from DINO-WM. DINO-WM is MIT-licensed; DINOv2 and MuJoCo use Apache-2.0. See THIRD_PARTY_NOTICES.md for exact provenance.

License

New probe code and documentation are released under the MIT License. Vendored components retain their upstream notices.