COUNTERMEM: World-Model Verified Counter-Factual Memory for Language Agents

Hongji Pu; Ruixiang Tang; Yongfeng Zhang

cs.AISep 25, 2026

COUNTERMEM: World-Model Verified Counter-Factual Memory for Language Agents

Authors: Hongji Pu, Ruixiang Tang, Yongfeng Zhang

Organizations: University of Illinois Urbana–Champaign · Department of Computer Science, Rutgers University

Abstract

Existing agent memory frameworks mainly create memory through an agent's interaction with the factual world, e.g., remembering feedback from actions taken to improve performance on future tasks. However, these frameworks seldom ask the "what if" question during memory construction: what if a different action had been taken, would the feedback have changed, and how could this feedback become useful memory? Obtaining such feedback directly in an active environment can be expensive and can alter the state needed for comparison. In this work, we introduce COUNTERMEM, a reinforcement-learning framework for constructing and using verified counterfactual memory across tasks. After a failed action, COUNTERMEM evaluates local alternatives from a copy or reset of the original state using executable world models, such as tests, proof checkers, and solvers. It stores improvements with the original and corrected actions, checked outcomes, and conditions for reuse. A learned memory-use policy selects a retrieved record or skips memory to balance task success and interaction cost, while the base LLM remains fixed. Both memory and policy are frozen during held-out evaluation. We evaluate COUNTERMEM on 12 benchmark settings across six domains. With gpt-oss-120b, COUNTERMEM improves both ReAct and Reflexion on all 12 benchmarks across six domains, averaging a gain of 12.6 percentage points over their unaugmented versions. In the four-domain comparison across two backbones, task-run tokens decrease by 7.7-42.0%, excluding offline selector-training costs. Further analyses show that removing verification or persistent storage weakens the gains, while applying verified corrections to unsuitable decisions can reverse them. Code will be released upon acceptance.

Figures & tables

Figure 1: Constructing and using counterfactual memory. Local alternatives are checked from a copy or reset of the original state. Verified improvements become records containing the action contrast, evidence, and applicability conditions. On later tasks, retrieval supplies compatible records, and a trained DQN selector chooses one or skips memory. Memory and selector parameters are frozen during held-out evaluation.

Base Agent	Augmentation	Math $\uparrow$	Coding $\uparrow$	Text-to-SQL $\uparrow$	SAT $\uparrow$	Tokens(M) $\downarrow$
Backbone: gpt-oss-120b
ReAct	None	29.8	48.7	26.8	50.7	6.0
	Best-of- $N$ (compute control)	34.5	38.1	36.2	41.1	9.4
	ExpeL ( Zhao et al., 2024 )	27.4	56.5	17.2	49.0	5.4
	Generative Agents ( Park et al., 2023 )	22.6	37.1	17.6	41.6	6.2
	Voyager ( Wang et al., 2024 )	21.9	60.6	26.3	51.4	5.2

Table 1: Four-domain comparison of memory mechanisms across agents and backbones. Scores (0–100, higher is better) average benchmarks within each domain; tokens report total task-run consumption in millions. Bold marks the best value within each agent–backbone group.

Figure 2: Learning from more source tasks and choosing what to keep. (a) Test score at successive memory-construction checkpoints. (b) Retrieval and context tokens per task; other solving costs are excluded. (c) Reported score loss on earlier tasks. (d) Test score under different record-retention rules at equal capacity. The horizontal axis in (a–c) counts source tasks, not stored records; shading shows one standard deviation across task orderings.

Figure 3: Reusing memory under four changes to the task. (a) Task success, with a no-memory reference. (b) Frequency of obtaining retrieved guidance. (c) Reported harmful-use percentages. “New surface” changes how the same error appears; “look-alike” preserves similarity while changing the condition needed for the correction. Definitions and reporting details are in Appendix D.4 .

Figure 4: Removing components and changing record assignment. (a) Score gains over the base agent. “CF without WM” stores unverified advice; “WM without CF” checks current-task alternatives without keeping cross-task corrections. (b) Evaluator calls per solved task. (c) Score gains when candidate records are absent, random, shuffled across states, or matched to the current decision. (d) Action changes and harmful changes per retrieved record.

Appendix figures & tables10 assets

Supplementary material from the paper’s appendix.

Appendix

Field	Content and purpose
$q$ : situation	A numeric threshold predicate may omit its boundary value.
$a^{-}$ : inferior action	The original query uses age > 18 .
$a^{+}$ : improved action	Replace that predicate with age >= 18 .
$e$ : evidence	The original and edited queries’ result-check outcomes, from the same database state.
$c$ : conditions	The question requests “at least” the threshold; the column and comparison types are compatible.
$u$ : reuse statistic	Initially zero; later summarizes observed helpful, neutral, or harmful reuse during construction.

Appendix

Table 2: Illustrative memory record for an inclusive age boundary.

Setting	Value
Training decision steps	2,000
Hidden layers	$64,64$
Activation	ReLU
Optimizer	Adam, $\epsilon_{\mathrm{Adam}}=10^{-8}$
Learning rate	$10^{-3}$
Learning starts	100 decisions

Appendix

Table 3: DQN configuration used for the reported memory-use selector.

Domain	Benchmark	Evaluator	Score	Decision unit
Math	miniF2F ( Zheng et al., 2022 )	Lean kernel	pass rate	tactic / proof step
Math	ProofNet ( Azerbayev et al., 2023 )	Lean kernel	pass rate	proof continuation
Math	AMBER ( Yang et al., 2026 )	Lean-based checker	pass rate	tactic / lemma choice
Coding	SWE-bench Verified ( Jimenez et al., 2024 ; Chowdhury et al., 2024 )	repository tests	resolved rate	edit / patch
Coding	BugsInPy ( Widyasari et al., 2020 )	regression tests	repair rate	localized edit
Formal	DafnyBench ( Loughridge et al., 2025 )	Dafny verifier	verification	invariant / annotation

Appendix

Table 4: Evaluation suite and executable checker used as the world model.

Condition	What changes	Example of the reuse question
Same environment, new task	A different request within the source environment.	Does the boundary correction help a new query over the same schema?
New environment, new task	The execution setting as well as the request.	Can a predicate correction apply in a different database with compatible types?
Same error, new surface	The form of the task or action, while the underlying mistake is similar.	Does the record match a renamed column or a differently worded inclusive request?
Look-alike, changed condition	Surface similarity remains, but a requirement for the correction fails.	Can the agent avoid changing $>$ to $\geq$ when the request requires a strict boundary?

Appendix

Table 5: Meaning of the four transfer conditions, with illustrative examples.

Figure 5: Adding CF Memory to each base agent on 12 benchmarks. The LLM is gpt-oss-120b. Each cell compares the augmented agent with its own unaugmented version. Positive entries mean higher scores in (a) or fewer evaluator calls and failed attempts per solved task in (b,c). Negative reductions mean additional interaction.

Benchmark	Score gain	Call reduction	Failure reduction
miniF2F	31.5 / 15.8	1.1 / 0.4	0.8 / 0.3
ProofNet	24.2 / 12.1	0.9 / 0.3	0.7 / 0.4
AMBER	27.4 / 9.9	2.4 / 0.9	1.6 / 0.6
SWE-bench Verified	21.4 / 8.4	3.6 / 1.2	1.1 / 0.5
BugsInPy	14.4 / 3.8	2.7 / 1.4	1.4 / 0.7
DafnyBench	11.2 / 5.2	0.6 / $-0.2$	0.3 / 0.1

Appendix

Table 6: Values in Figure 5 . Each entry is ReAct / Reflexion. Score changes are pp; call and failure reductions are per solved instance.

Template	Used in	Required behavior
GENERATE_LOCAL ALTERNATIVES	Module I	Return local edits in the fixed edit grammar, one changed decision unit per candidate; do not generate a complete replacement submission.
DISTIL_RECORD	Module II	Distill the checked contrast into a reusable record. The condition must be checkable on a future task and must be able to be false on a look-alike task.
DRAFT_ACTION	Module III	Produce the action draft before any memory is shown. This is the memory-free draft used by the selector and reuse-credit logic.
REVISE_WITH_RECORD	Module III	Show exactly one selected record and revise the draft only when useful; the model may return the draft unchanged.
RECORD_RENDER / RECORD_RENDER_NO CONDITION	Condition ablation	Render the same record fields in both cases; the only intended difference is whether the applicability clause is shown.
BEST_OF_N_JUDGE	Best-of- $N$ control	Choose exactly one candidate and return CHOICE: <n> . Parse failures are logged and fall back to candidate 1; the verifier does not choose the candidate.

Appendix

Table 7: Prompt interfaces used by the reported runs.

Figure 6: Four-domain summaries of the main table. Each arrow links one unaugmented agent to its CF Memory version. The right panel plots gain against base score; the shaded band is variation across the four agent–backbone combinations, not a confidence interval over runs.

Figure 7: Comparison at a common 6M-token budget. CF Memory outperforms ReasoningBank, ReasoningBank with MaTTS, and ReasoningBank with verified feedback in all five domains. Reported averages give each domain equal weight.

Figure 8: Benefits and harm from reuse. (a) In 120 paired runs per domain, rescued tasks fail without memory and succeed with it; broken tasks have the opposite outcome. Diamonds mark the net count, rescued minus broken. (b) Removing only the applicability field

c

increases harmful use, especially on look-alike tasks.

Explore similar work

Beyond the Shadows of Plato's Cave: Evaluating False Memory in Autonomous Agents via Counterfactual Reasoning
Sep 30, 2026Quan M. Tran, Zhuo Huang, Zhen Fang +3Agent Memory
Sydney AI Centre, The University of Sydney · Australian Institute for Machine Learning, Adelaide University · Australian Artificial Intelligence Institute, University of Technology Sydney +2
From Attack Success to Attack Severity: Counterfactual Memory Attacks on LLM Agents
Sep 28, 2026Mingxi Zou, Langzhang Liang, Zhuo Wang +3Adversarial Attacks Agent Memory Poisoning
Shanghai Academy of AI for Science (SAIS) · Fudan University · Monash University
Many Steps into One: Reasoning Contraction and the Limits of Safe Memory Reuse for LLM Agents
Oct 13, 2026Jianghong Zhou, Hao Rong, Chen Lin +1Agent Memory Management LLM Agent Safety
ByteDance · Scale AI

Sep 30, 2026J/K move · Enter open · S save

Beyond the Shadows of Plato's Cave: Evaluating False Memory in Autonomous Agents via Counterfactual Reasoning

Quan M. Tran, Zhuo Huang, Zhen Fang, Jing Zhang+2

Sydney AI Centre, The University of Sydney · Australian Institute for Machine Learning, Adelaide University · Australian Artificial Intelligence Institute, University of Technology Sydney+2

Agent Memory

Open paper AskPDF arXiv

Autonomous agents increasingly rely on memory to generalize beyond their training environments. However, agents are bounded by what they have seen and believed, and leveraging such memories in unseen environments can introduce biases into their internal beliefs. We formalize this phenomenon as false memory, which can arise from spurious correlations, environment shifts, and knowledge conflicts. Despite its importance, false memory is difficult to evaluate because it stems from agent internal beliefs and is easily confounded with ordinary generalization failures. Therefore, we propose FAME, a training-free framework that evaluates false memory through the evolution of agent beliefs under counterfactual reasoning. Specifically, counterfactual scenarios reveal how beliefs change as the latent concept of memory shifts under hypothetical interventions; thus, measuring the resulting concept drift provides a signal for distinguishing faithful versus false memory. Such concepts can be estimated from agent hidden states before answer generation, avoiding the need for reward design or answer sampling. Empirical experiments reveal that simply monitoring answers often fails to detect false memory, while FAME achieves AUROCs of 76.2% - 96.7% across false-memory settings, and outperforms the best baseline by 3.4% - 23.3% across realistic benchmarks, spanning math reasoning (GSM-Symbolic), code generation (GitChameleon), and complex reasoning (BigBench-Hard). We further release corresponding counterfactual templates and facilitate future research on false memory.

Abstract

Explore similar work

Beyond the Shadows of Plato's Cave: Evaluating False Memory in Autonomous Agents via Counterfactual Reasoning

From Attack Success to Attack Severity: Counterfactual Memory Attacks on LLM Agents

Many Steps into One: Reasoning Contraction and the Limits of Safe Memory Reuse for LLM Agents