PRIME

Reference implementation for the paper Principled Direction-Free Intrinsic Motivation through Model-Free Epistemic Free-Energy Estimators.

PRIME is a single window-stationary intrinsic reward derived from parameter information gain. It produces surprise-maximizing behavior where dynamics are unresolved and surprise-minimizing behavior where they are resolved, under one configuration across two didactic environments: Butterflies (high-entropy) and Maze (low-entropy).

This was initially a passion project between two friends, but has never reached the level of maturity we hoped for. I have decided to instead go for a minimal "proof of concept" to demonstrate at least the core idea which motivated the work in the first place.

Unlike my other projects, getting started and reproducing everything is a lot more straight forward, because it doesn't involve preparing any datasets. However, if you have any inquiries or wish to discuss the work, contact me at a.furutanpey@coovally.ai.

Setup

./setup.sh

Creates .venv and installs the pinned dependencies. Requires Python 3.12.

Run

./run.sh [method] [env] [seed]

Defaults to prime maze 0. Logs are written to logs/<method>_<env>_seed<seed>/.

method	description
prime	PRIME, canonical config (K=5, α=0.5)
prime_k1	ablation, single statistics head (K=1)
prime_gateoff	ablation, look-ahead gate off (α=0)
rnd	Random Network Distillation
smirl	Surprise Minimizing RL
disagreement	ensemble disagreement
extrinsic	task-reward DQN (supervised reference)
random	analytical random walker

env is butterflies or seed is any integer. Example:

./run.sh prime butterflies 0
./run.sh rnd maze 3

A run trains for the per-seed budget in the config (250k env-steps on Butterflies, 200k on Maze) and selects cuda, then mps, then cpu automatically.

Results

.venv/bin/python -m analysis.report logs/

Prints per-run task reward: mean, peak rolling-20, and last rolling-20. On Butterflies this is catches per episode; on Maze it is the goal-reach rate. Training curves are also viewable with tensorboard --logdir logs/.

Full reproduction

./sweep.sh

Runs every method on both environments for seeds 0–14, the grid behind Table 1 and Table 2.

Layout

prime/        the method: reward (Eq. 10), estimators, ensembles, window freeze
envs/         Butterflies and Maze
baselines/    RND, SMIRL, disagreement, extrinsic-DQN, random walker
configs/      one YAML per method and environment
analysis/     read run logs into headline metrics

Comments in the source point to the corresponding equations and sections of the paper.