Skip to content

Repository files navigation

nanoAWM

nanoAWM consequence-simulation diagram

A tiny, inspectable Agent World Model lab for long-horizon tool agents.
Local symbolic evidence in a deterministic environment — not external validation, not real-agent safety.

Python License DOI Tests Scope External validation

Agents should dream before they act.

nanoAWM trains compact, action-conditioned world models that predict the consequences of an action before it runs (next observation, state diff, reward, risk, reversibility, approval violation, hidden-state corruption, and termination), then plans against those predictions. It is meant to be small enough to read in an afternoon and honest enough to attack.

It is not a chatbot, a framework, a benchmark-only repo, or a claim about production agent safety. It is a research toy in the spirit of nanoGPT: one idea, made concrete, with its limitations in plain sight.

The idea: consequence aliasing

Long-horizon tool agents fail when they act from the visible transcript alone. Two actions can look equally reasonable from the current observation while diverging sharply in future reward, reversibility, hidden-state corruption, or approval compliance. We call this consequence aliasing.

flowchart LR
    O["same visible observation"] --> A1["action A: locally plausible"]
    O --> A2["action B: locally plausible"]
    H["hidden state / policy / delayed tests"] -.-> A1
    H -.-> A2
    A1 --> C1["recoverable success"]
    A2 --> C2["irreversible failure"]
    WM["action-conditioned world model"] --> P["imagined consequences"]
    P --> PL["risk-aware planner"]
    PL --> A1
Loading

The narrow, falsifiable question this repo asks:

In a deterministic symbolic OS with hidden state and delayed consequences, can a learned consequence model help a planner avoid irreversible / approval-violating / corrupting actions better than reactive and rule-based baselines, without seeing labels or oracle state at planning time?

Inside this lab, the answer is yes, but read the honest results and what is contaminated before you believe the big number.

Results (honest version)

All numbers are from the local MiniOS suite (420 tasks, 43,164 transitions; the main model is 27,684 trainable head parameters trained on the 24,694-transition train split). Regenerate with scripts/repro_core.sh.

Held-out task success — the comparison that matters:

Policy Held-out success Notes
learned_world_model_planner 1.000 plans against predicted consequences
oracle_world_model_planner 1.000 upper bound; sees true transitions (marked contaminated)
best non-world-model baseline 0.867 scripted_expert / heuristic_llm_risk_guard
rule_based_risk_guard 0.800 hand-written safety rules
reactive_policy 0.067 adversarial baseline — by construction always takes the action labelled risky

The honest headline is learned − best baseline = +0.133, not +0.933. The +0.933 gap over reactive_policy is real within this benchmark but reactive_policy is defined to lose (reactive_action is set to the risky action for every task), so it is a floor, not a fair opponent. We report it labelled as such rather than on a badge.

One-step prediction (test split, 6,170 transitions):

Metric Value
Next-observation accuracy 0.996
State-delta accuracy 0.987
State-delta macro-F1 0.791
Risk AUROC ~1.000
Expected calibration error 0.029

(Earlier versions of this repo reported "state-delta F1 = 0.980"; that number was accuracy mislabeled as F1. The true multiclass macro-F1, which averages per-class scores and so penalizes the rare state-diff classes, is 0.791. Fixed.)

Genuinely held-out (the honest number)

Is the 1.000 real? No — it is in-distribution (eval shares training vocabulary; see the next section). The honest test is to make the model earn success on a vocabulary it has never seen. heldout_eval retrains the model with vocabulary randomization (every object name and action marker is replaced by fresh random tokens on every training variant, so nothing is memorizable) and evaluates on a suite whose object names and action markers are disjoint from training (verified token-set disjointness). Regenerate with python -m nanoawm heldout_eval (reports/heldout_scorecard.md).

Policy Genuinely held-out success (105 tasks)
learned_world_model_planner 0.524
best non-world-model baseline 0.067
oracle_world_model_planner 1.000 (upper bound)

Honest improvement over the best baseline on genuinely-held-out tasks: +0.457. The in-distribution 1.000 is a sanity check; this 0.524 is the number to believe.

The per-family breakdown is the real result — it reveals a clean law:

The world model transfers to novel vocabulary exactly where the environment grounds the safe/risky distinction in an observation, and only fails where the distinction is encoded in arbitrary action vocabulary with no observable grounding.

Where grounding is observable, held-out success is high (package_install 1.000, stale_config 1.000, irreversible_file_deletion 1.000, rollback/database/cross_surface_dependency 0.571 — distinct from the harder browser_terminal cross-surface OOD suite below). Where the safe and risky actions differ only by an arbitrary marker the clue never disambiguates (approval_gated_email_send), held-out success is 0.000, and no consequence model could recover that without memorizing the marker. That is a precise, falsifiable boundary, not a number to hide.

The grounding is learned, not hand-coded. A within-scorecard ablation (same randomized training) shows the learned grounding feature, which emits the observed clue words adjacent to each candidate's tokens and lets the model learn which context words mean "safe", beats a simple token-overlap grounding feature 0.524 vs 0.419, while staying strictly inside the anti-cheat boundary (it never surfaces a label-name token).

What is, and is not, shown here

This is the section a skeptical reviewer should read first.

The environment is closed, deterministic, and partly lexically separable.

  • The eval splits are drawn from the same small vocabulary as training, so they overlap training at the feature level the model actually consumes: test 90%, template_holdout 100%, ood 71% feature-identical to a training task (measured and reported by anti_cheat; see reports/anti_cheat.jsoncontamination). The split named template_holdout does not hold out templates. Treat "held-out 1.000" as in-distribution performance: the genuinely-held-out scorecard above (disjoint vocabulary) is 0.524, and that is the honest generalization number.
  • The safe/risky distinction is partly encoded in the action strings (validated_change vs direct_change, rm -rf, …). A pure good/bad-keyword matcher with no world model scores ~0.73. So a large part of the task is lexical, not consequence reasoning.

Because of that, the most credible evidence is not the 1.000. It is the genuinely-held-out scorecard (0.524) above, the generalization ladder, the mechanism audit, and the negative results:

  • Mechanism (consequence aliasing audit): over 15,736 same-visible-state candidate-action pairs, the learned model ranks the safer hidden consequence higher with choice accuracy 0.971 and risk-order accuracy 1.000, using scorer-only counterfactuals after candidates are fixed (reports/aliasing_audit.md). This is the closest thing to direct evidence for the central claim.
  • Negative result — action paraphrase: when action target/value strings are replaced by unseen aliases (hidden roles preserved), learned success collapses while the oracle stays at 1.000. The model leans on action vocabulary; we do not claim action-vocabulary robustness (reports/action_paraphrase_audit.md).
  • Negative result — cross-surface OOD: on a 180-task suite requiring both a browser probe and a terminal canary before the safe action, the oracle solves it (1.000) but the learned planner gets 0.011 and the calibrated planner 0.017. A genuine, unrecovered failure, kept out of the main scorecard (reports/browser_terminal_ood.md).

These failures are the point. They show where a compact consequence model stops working, which is more useful than another saturated metric.

Generalization ladder (genuinely held-out)

The 1.000 headline is in-distribution (the eval splits overlap training vocabulary, above). To measure what the model actually generalizes, the ladder rebuilds the suite over a disjoint vocabulary (held-out object names and disjoint nonce action markers that share no tokens with training) and asks a single controlled question: with the vocabulary held fixed-novel, what kind of observation grounding lets the model transfer? Regenerate with python -m nanoawm generalization_ladder --run-ladder (reports/generalization_ladder.md).

Two separate axes — do not read them as one ramp.

Invariance check (the shipped model, one manipulation):

Rung Condition Learned success
A disjoint object names, shared action markers 1.000

Rung A only asks whether the committed model is invariant to renaming objects. It is. It depends on markers, not object names, which is exactly what the next axis interrogates.

Grounding ladder (one held-out suite, one training protocol, fresh matched models, 3 seeds each):

Rung Held-out condition Learned success (mean) Per-seed
B disjoint markers, no grounding 0.350 0.28 / 0.36 / 0.42
C disjoint markers, direct grounding (clue names the safe action) 1.000 1.0 / 1.0 / 1.0
D disjoint markers, indirect relational grounding (clue states a property; model must infer) 1.000 1.0 / 1.0 / 1.0

B, C and D are evaluated on the same 98 held-out tasks with the same training protocol (fresh models on the matching disjoint-marker train split, identical epochs), across seeds 71/73/77. The only thing that changes is how the safe action is grounded in the probed observation. So the rungs are directly comparable:

  • B — the real limit. With novel action markers and no grounding cue, transfer collapses to a 0.350 floor (residual non-marker signal only). Without grounding, the model cannot identify the safe action under unseen vocabulary. This is the same dependence the action-paraphrase negative result shows on a different held-out construction.
  • C — direct grounding recovers it (1.000). When the observation names the safe action's own markers, the planner follows the pointer to novel vocabulary. Pointer-following.
  • D — indirect grounding also recovers it (1.000), and this is the harder mechanism. The clue never names which action is safe; it states a property — a relation word (active vs stale) adjacent to each candidate, extracted by a hand-coded feature (see boundary #1 below) — and the model must learn the property → safety mapping (not what the words mean) to act. Pointer-following is structurally impossible here: the direct-grounding features are identical for the safe and risky candidates (both values appear in the clue, neither target does), so the relation polarity is the only feature that separates them (locked by tests/test_generalization_ladder.py). The polarity is also "none" until the planner takes a probe action — the signal is earned by acting, not handed over at t0.

The controls are the evidence, not the 1.000. Two independent control families, both averaged over the three seeds:

Control C (direct) D (relational) What it shows
shuffled grounding (point the cue at the risky action) 0.000 0.000 causal: nothing changes but which marker the cue names, and success goes 1.000 → 0.000 — the model follows the grounding into total failure
grounding model on no-grounding eval 0.000 0.000 confound: the grounding-trained models collapse when grounding is absent, so the recovery is grounding-driven, not a stronger model

Honest boundaries. (1) The relation lexicon (active/stale/verified/deprecated/…) is shared between train and eval, and rung D's extraction is hand-coded (_relational_polarity parses "active <marker>" into a polarity feature). What the model learns and transfers is the polarity → safety mapping applied to disjoint object/action vocabulary; it is not learning to parse language or generalizing to an unseen relation lexicon. The lexicon-holdout experiment (python -m nanoawm lexicon_holdout) measures exactly this dependence with a learned extractor instead (the learned_grounding feature, not hand-coded): trained on rotating relation pairs, it holds 1.000 on a seen pair under disjoint object/action vocabulary but collapses to 0.582 on an unseen pair (sound/suspect, never in training) — a −0.418 drop on the same tasks and model, with clause-order randomization controlling the positional-cue confound (locked by tests/test_lexicon_holdout.py). Generalizing the lexicon itself needs semantic priors this lab deliberately omits. (2) C and D both saturate at 1.000 because, once grounding is present and read, the safe action is cleanly separable in this deterministic 98-task suite — the defensible claim is the floor → recovery contrast and the causal/confound controls, not the saturation. (3) This is local symbolic evidence, not external validation.

MiniOS

A deterministic symbolic computer-use environment with enough surface area to expose consequence aliasing without browser-automation noise: filesystem reads/writes/rollback/irreversible-delete, a test runner with delayed failures, a git workflow, browser forms with hidden validation, database mutations with hidden invariants, email/calendar approval gates, and package/config changes with dependency conflicts. Hidden-state digests are logged for the scorer but never exposed to learned planners.

Representative actions:

READ_FILE, WRITE_FILE, RUN_TESTS, GIT_STATUS, GIT_COMMIT,
OPEN_PAGE, CLICK, FILL_FORM, SUBMIT_FORM, QUERY_DB, UPDATE_DB,
DRAFT_EMAIL, SEND_EMAIL_SIMULATED, CREATE_EVENT_SIMULATED,
INSTALL_PACKAGE_SIMULATED, REQUEST_APPROVAL, ROLLBACK, NOOP

15 task families span safe/stale config updates, irreversible deletion, approval-gated send, hidden DB invariants, hidden browser validation, dependency conflicts, dirty git trees, calendar/email conflict, delayed test failure, rollback recovery, ambiguous same-observation actions, multi-step repair, cross-surface dependencies, and adversarial variants.

Models

  • Main world model — a from-scratch tiny recurrent multi-head predictor (tiny_recurrent_multihead_v2) over signed-hash symbolic features, trained with pure-Python SGD. Given history, visible state, observation, and a candidate action it predicts next-observation class, state-diff label, reward, risk, reversibility, approval violation, hidden-state corruption, and done.
  • Dense sequence probe (tiny_dense_sequence_multihead_v1) — a no-dependency dense sequence model for the same targets, kept as an inspectable point of comparison before any framework-backed model is introduced.
  • Transformer probe (tiny_transformer_world_model_v1) — a genuine from-scratch transformer (learned token + positional embeddings, single-head self-attention, FFN, two layernorms, residuals, mean-pool, the same multi-head outputs) built on a micrograd-style scalar autodiff engine (nanoawm/autograd.py) so it is dependency-free and gradient-correct (finite-difference checked; it cleanly overfits a tiny batch). The scalar engine is slow, so it is a small sampled probe, not the main planner. Honest comparison (python -m nanoawm transformer_eval, same 300 test rows): the transformer reaches one-step next-state 0.773 / state-delta 0.767 / risk-AUROC 0.943 vs the hashed model's 0.997 / 0.993 / 1.000 — a learned-embedding model does not beat the hashed model here, and on the disjoint-vocabulary suite it cannot (held-out tokens are out-of-vocabulary). "Scale the model" needs semantic pretraining this lab deliberately does not have; the probe shows the model class is supported and measures, rather than assumes, the (null) benefit.

Planners and baselines

Required baselines: reactive_policy, plan_only_no_world_model, rule_based_risk_guard, oracle_world_model_planner, learned_world_model_planner.

Stronger local baselines (no model inference, just heuristic routing over public text/state, named accordingly): scripted_expert_policy, heuristic_planner_policy, heuristic_llm_style_policy, heuristic_llm_risk_guard_policy, and a model-assisted heuristic_llm_world_model_policy. (Genuinely local open-weight LLM baselines via Ollama / LM Studio are optional and kept behind acclaim_evidence; they are weak baselines, not external validation.)

The learned planner uses candidate actions, one- and multi-step rollouts, risk/uncertainty penalties, approval constraints, irreversible-action penalties, and receding-horizon execution. It never sees real next state before choosing.

Anti-cheat boundary

At planning time the learned planner does not receive: task/template/split ids, safe/risky/reactive labels, rewards or risk as features, hidden gold state, oracle next state, or post-hoc scorer labels (nanoawm/features.py enforces the feature allow-list; nanoawm/anti_cheat.py checks it and reports train/eval contamination honestly). Oracle baselines are explicitly flagged oracle_contaminated=true. The scorer may use hidden fields only post-hoc.

Quickstart

No paid APIs, no network. Reproduce the core result from scratch:

bash scripts/repro_core.sh

That runs the eight-command core loop (generate → collect → train → eval → run → score → anti-cheat → aliasing audit) and writes reports/world_eval.json, reports/anti_cheat.json, reports/aliasing_audit.md, and a fast-smoke reports/scorecard.smoke.json (compare its headline fields to the committed full reports/scorecard.json). The smoke run uses --skip-expensive, so it deliberately does not regenerate the expensive audit reports — run python -m nanoawm.score without --skip-expensive to regenerate those (strong baselines, scaling sweep, exploitation, etc.).

Run the tests:

uv run --with pytest python -m pytest -q

Optional deeper audits each have their own subcommand (python -m nanoawm --help): heldout_eval (the genuinely-held-out disjoint-vocabulary scorecard), generalization_ladder (--run-ladder trains the grounding rungs B/C/D), lexicon_holdout (seen-vs-unseen relation lexicon, the final mapped boundary), aliasing_audit, adversarial_ood, action_paraphrase_audit, grounded_paraphrase_audit, compositional_grounding_audit, calibrated_safe_planner, model_exploitation_audit, browser_terminal_ood, browser_terminal_grounding, browser_terminal_retraining, plus browser_bench (Playwright) and acclaim_evidence (scaling sweep, statistics, optional local-LLM baselines).

data/trajectories.jsonl and the larger training-trace files are generated, not committed (they can exceed GitHub's file-size limit); recreate them with the commands above.

Reading map

What this is not claiming

This repo does not claim that symbolic MiniOS performance transfers to real-world autonomous agents, robustness under arbitrary distribution shift, or production safety. It does not claim external validation, public reception, or critical acclaim — there are currently zero independent reproductions, expert reviews, or citations.

The claim is narrower and defensible:

A compact learned consequence model can improve tool-agent planning in a controlled environment with hidden state, delayed consequences, approval boundaries, rollback, and irreversible actions, while preserving anti-cheat separation at planning time — and it has measurable, documented failure modes (action-vocabulary dependence, cross-surface OOD).

Limitations and next steps

  • "Held-out" now means something. The genuinely-held-out scorecard (disjoint object names and action markers, vocabulary-randomized training) is the honest headline: learned planner 0.524 vs 0.067 baseline (+0.457), with a per-family law that the model transfers exactly where grounding is observable. The generalization ladder is the controlled sweep behind it: no grounding → 0.350 floor, grounding → recovery, shuffled controls → 0.000. Remaining work: the families that stay at 0.000 (e.g. approval_gated) encode safe/risky in arbitrary vocabulary with no observable grounding, so closing those needs richer task grounding or a model with semantic priors, not just more data.
  • Beyond lexical cues. The genuinely-held-out 0.524 (vs in-distribution 1.000) is the precise size of the lexical-dependence gap. The learned grounding feature recovers it where the observation grounds the distinction; the residual is families with no observable grounding and the action-paraphrase / cross-surface OOD negatives.
  • The relation lexicon must be shared — measured, not just asserted. The lexicon-holdout experiment trains the learned-grounding model on rotating relation pairs and evaluates the same tasks with a seen pair vs an unseen one: success holds at 1.000 on the seen pair under disjoint object/action vocabulary but collapses to 0.582 on the unseen pair (sound/suspect) — a −0.418 drop, controlled for the positional-cue confound by clause-order randomization (tests/test_lexicon_holdout.py). This is the precise edge of generalization: the model transfers polarity over novel objects and actions, but cannot infer the polarity of a relation word it has never seen. Closing it needs semantic priors (pretraining), not more symbolic data — the same conclusion the transformer probe reaches from the other direction.
  • Independent reproduction. Two clean-clone reproductions by independent reviewers before any stronger public claim.
  • Scale the model — done, and it doesn't help yet. A from-scratch transformer probe (tiny_transformer_world_model_v1, pure-Python autodiff) now exists and is gradient-correct, but it underperforms the hashed model on a sample (0.773 vs 0.997 one-step) and cannot help on disjoint vocabulary (OOV tokens). The honest next step is semantic pretraining (so unseen vocabulary has a representation), not a bigger from-scratch model on this tiny data.

Personal project disclaimer

Personal, independent research and development. Not affiliated with, endorsed by, or representative of any employer, client, or organization.

About

Tiny, inspectable Agent World Model lab for learned consequence simulation in tool agents.

Topics

Resources

Contributing

Security policy

Stars

0 stars

Watchers

0 watching

Forks

Releases

Packages

Contributors

Languages