Visualizing the World Model in an Agent's Value Function

Shawheen Naderi; Fatemeh Akherati; Yubei Chen

cs.AIOct 9, 2026

Visualizing the World Model in an Agent's Value Function

Authors: Shawheen Naderi, Fatemeh Akherati, Yubei Chen

Organizations: University of California, Davis

Abstract

Could model-free reinforcement learning algorithms actually be learning a model of the environment? In this work, we apply counterfactual interventions which hold everything about a state fixed, besides a targeted change, to explore if the value function of a Proximal Policy Optimization agent forms a world model of the Crafter environment. Despite only being trained to optimize reward, we discover that the value function learns concepts such as the rules governing when a diamond is collectible, water and food having an impact on health, heuristics regarding inventory, and more. Finally, we discuss the safety implications of our visualizations and some shortcomings of current reinforcement learning agents that will prevent us from building truly advanced machine intelligence.

Figures & tables

Panel of Figure 1 — Figure 1: In panel (a), the pair shows the agent before and after crafting a stone pickaxe which is highlighted by the yellow box, and a cow also appears highlighted by the red box. Crafting the pickaxe requires the agent to be near a table and possess the necessary wood and stone resources; afterward, the stone pickaxe appears in the inventory, enabling the agent to collect iron. The external environment is composed of 18 possible textures, while the internal environment contains 16 possible textures including inventory items and health-related components. The agent’s view of the environment is only partially observable. Panel (b) shows the achievement graph, where arrows indicate prerequisite relationships between achievements.

Figure 2: The curve shows the raw predictions, the blue curve shows the smoothed value estimate, and the markers indicate achievement, health, and negative rewards. The predicted value declines rapidly as achievements are collected and remains near zero for most of the later episode, with a brief increase near the end when a diamond is collected.

Figure 3: Starting with the earliest steps on the left, we see that Crafty has an empty inventory and has yet to collect any resources. As Crafty continues through the episode, it starts expanding its inventory which allows it to enter cave areas, but also expends the finite rewards in the environment. By step 300, the agent has collected almost all the achievements, excluding diamond, and thus has a very low predicted value until it finds the diamond at step 610.

Figure 4: Starting from an unaltered state with an empty inventory, we sequentially add wood, a wood pickaxe, stone, and a stone pickaxe while keeping the rest of the state fixed. The critic’s predicted value generally decreases as the inventory progresses.

Panel of Figure 5 — Figure 5: Achievement success rates with and without the penalty.

Panel of Figure 6 — Figure 6: Three different counterfactual cases for diamond, iron, and plant. Changes are denoted with red and yellow bounding boxes.

Panel of Figure 7 — Figure 7: The blue curve shows the smoothed success rate at which the agent collects the achievement during training and is defined by the percentage of episodes in which the agent successfully collected the achievement at least once. The red curve shows the fraction of fixed counterfactual states for which the critic satisfies the expected ordering.

Figure 8: $\Delta_{R}(h,r)$ is a column’s bottom row’s value minus it’s top row’s value. We observe that the delta between having water in the environment is larger when internal water levels are low compared to when they are high indicating that the agent values water more when it has a larger impact on its health. Further, the spread is enhanced when the agent the agent is low on health and the risk of death due to dehydration is higher.

Appendix figures & tables10 assets

Supplementary material from the paper’s appendix.

Appendix

Hyperparameter	Value
Environment	Craftax-Classic-Symbolic-v1
Total environment interactions	$10^{10}$
Parallel environments	1,024
Rollout length	64
Interactions per update	65,536
PPO update epochs	4

Appendix

Table 1: PPO training configuration.

Achievement	$10^{5}$	$10^{6}$	$10^{7}$	$10^{8}$	$10^{9}$	$10^{10}$	$10^{11}$
Collect Coal	0.0	0.2	28.6	85.6	95.9	100.0	100
Collect Diamond	0.0	0.0	0.0	0.7	14.9	30.5	41.7
Collect Drink	7.0	25.1	60.9	85.0	95.0	100.0	97.2
Collect Iron	0.0	0.0	0.9	76.5	93.4	98.3	100
Collect Sapling	33.9	93.7	97.7	97.4	88.4	98.3	100
Collect Stone	0.0	1.6	76.8	97.4	100.0	100.0	100

Appendix

Table 2: Achievement success rates for the trained Crafty agent. The success rate for an achievement is the percentage of evaluation episodes in which the agent successfully completed that achievement at least once. Tasks lower on the achievement hierarchy such as collect wood are quickly learned by the agent but more complicated achievements require more environment steps to be learned. Note: PPO logs metrics once every 65,536 environment steps. To show learning across orders of magnitude, each column labeled

10^{n}

reports the nearest available W&B record: 131K, 983K, 10.027M, 100.008M, 1.000B, 9.999B and 100.000B steps, respectively.

Panel of Figure 11 — Figure 11: Representative counterfactual inventory swaps between early- and late-episode observations. The three examples are placed side by side for direct comparison.

Panel of Figure 12 — Figure 12: Representative diamond-collectability counterfactuals. Each example displays the four combinations of diamond presence and iron-pickaxe possession, with all remaining state features held fixed.

Panel of Figure 13 — Figure 13: Representative iron-collectability counterfactuals. Each example displays the four combinations of iron presence and stone-pickaxe possession, with all remaining state features held fixed.

Quantity	Water	Cow/food
Eligible states	9,473	9,692
Primary success rate	82.25% [81.48, 83.01]	83.24% [82.51, 83.98]
$\Delta_{R}(2,0)>0$	94.46%	99.49%
$\Delta_{R}(2,0)>\Delta_{R}(2,9)$	98.47%	99.82%
$\Delta_{R}(2,0)>\Delta_{R}(9,0)$	85.63%	83.53%

Appendix

Table 3: Health-dependent resource-valuation results. Confidence intervals are 95% percentile intervals from 10,000 paired episode-bootstrap resamples.

Panel of Figure 14 — Figure 14: Representative water counterfactuals. Rows vary visible-water presence, while columns vary health and internal drink level.

Panel of Figure 15 — Figure 15: Representative food counterfactuals. Rows vary visible-cow presence, while columns vary health and internal food level.

World Models

World models represent and predict how environments change, helping agents plan, act, and reason about possible futures.

See all 366 papers

Learning Implicit Causal World Models from Multi-Agent Demonstrations
Jul 28, 2026Jasorsi GhoshCausal Representation Learning Causal World Models
Department of Computer Science Purdue University West Lafayette, Indiana, USA
Inverting the Bellman Equation: From $Q$ -Values to World Models
Jun 19, 2026Alistair Letcher, Mattie Fellows, Alexander D. Goldie +3Reinforcement Learning Model-Free RL
FLAIR, University of Oxford · Google DeepMind · Mila, University of Montreal
What Do World Models Learn in RL? Probing Latent Representations in Learned Environment Simulators
Mar 23, 2026Xinyu ZhangWorld Models World Model Learning
Anyscale

Jul 28, 2026J/K move · Enter open · S save

Learning Implicit Causal World Models from Multi-Agent Demonstrations

Jasorsi Ghosh

Department of Computer Science Purdue University West Lafayette, Indiana, USA

Causal Representation Learning Causal World Models World Models Causal Structure Learning Causal RL Offline MARL+1

Open paper AskPDF arXiv

In model-based reinforcement learning, world models exist as internal simulators, but their training often conflates statistical correlations with causal mechanisms. This problem is exacerbated in multi-agent systems where physical transitions are intertwined with strategic agent intents, causing world models to fail under distribution shift. We introduce Implicit Causal World Models to recover environmental dynamics from offline demonstrations without requiring pre-defined causal graphs. By incorporating policy variance, we render world models discoverable via the sequential backdoor condition. Evaluations across coordination tasks (Two-Door, Navigation, and Giveway) demonstrate that these models provide interpretable causal representations under both full and partial observability, with model accuracy scaling directly with interventional strength.

Abstract

World Models

Explore similar work

Learning Implicit Causal World Models from Multi-Agent Demonstrations

Inverting the Bellman Equation: From QQQ-Values to World Models

What Do World Models Learn in RL? Probing Latent Representations in Learned Environment Simulators

Inverting the Bellman Equation: From $Q$ -Values to World Models