Reward Prediction with Factorized World States

Yijun Shen; Delong Chen; Xianming Hu; Jiaming Mi; Hongbo Zhao; Kai Zhang; Pascale Fung

cs.CLMar 10, 2026

Reward Prediction with Factorized World States

Authors: Yijun Shen, Delong Chen, Xianming Hu, Jiaming Mi, Hongbo Zhao, Kai Zhang, Pascale Fung

Organizations: East China Normal University · HKUST

Abstract

Agents must infer action outcomes and select actions that maximize a reward signal indicating how close the goal is to being reached. Supervised learning of reward models could introduce biases inherent to training data, limiting generalization to novel goals and environments. In this paper, we investigate whether well-defined world state representations alone can enable accurate reward prediction across domains. To address this, we introduce StateFactory, a factorized representation method that transforms unstructured observations into a hierarchical object-attribute structure using language models. This structured representation allows rewards to be estimated naturally as the semantic similarity between the current state and the goal state under hierarchical constraint. Overall, the compact representation structure induced by StateFactory enables strong reward generalization capabilities. We evaluate on RewardPrediction, a new benchmark dataset spanning five diverse domains and comprising 2,454 unique action-observation trajectories with step-wise ground-truth rewards. Our method shows promising zero-shot results against both VLWM-critic and LLM-as-a-Judge reward models, achieving 60% and 8% lower EPIC distance, respectively. Furthermore, this superior reward quality successfully translates into improved agent planning performance, yielding success rate gains of +21.64% on AlfWorld and +12.40% on ScienceWorld over reactive system-1 policies and enhancing system-2 agent planning. Project Page: https://statefactory.github.io

Figures & tables

Figure 1: RewardPrediction benchmark Overview. Given a textual goal description, the model computes step-wise progress estimates $[\hat{r}_{t}]_{t=0}^{n}$ from sequences of action-observation pairs. The reward predictions are compared against ground-truth $[r_{t}]_{t=0}^{n}$ using EPIC distance ( Gleave et al., 2020 ) to quantify alignment.

Figure 2: RewardPrediction Benchmark Overview. Representative trajectories across five diverse domains. Each column displays a task instance with actions, observations, and the corresponding ground-truth task progress score ( $R\in[0,1]$ ) at key time steps.

Figure 3: Two paradigms of reward prediction. Unlike representation-free methods that regress rewards directly from raw inputs (top), representation-based frameworks derive progress signals by measuring the alignment between factorized state representations $s_{t}$ and goal interpretations $g_{t}$ (bottom).

Figure 4: StateFactory Framework. State extraction and goal interpretation are coupled recurrent processes (left). Our approach factorizes states into explicit objects and attributes (right), deriving dense rewards from semantic similarity between $\hat{s}_{t}$ and $\hat{g}_{t}$ .

Method	Backbone	Training Data	Reasoning	Zero-shot	Reward Prediction Error ( $D_{\textsc{EPIC}}\downarrow$ )
Method	Backbone	Training Data	Reasoning	Zero-shot	AlfWorld	ScienceWorld	WebShop	BlocksWorld	TextWorld	Overall
Monotonic Baseline	–	–	–	✓	0.532	0.508	0.535	0.589	0.536	0.540
Supervised RM	Qwen2.5-1.5B	AlfWorld	N/A	✗	0.212	0.542	0.618	0.596	0.414	0.476
		ScienceWorld			0.506	0.305	0.661	0.706	0.580	0.552
		WebShop			0.707	0.624	0.242	0.706	0.707	0.597
		BlocksWorld			0.556	0.471	0.596	0.472	0.556	0.530

Table 1: Main Results: We compare StateFactory with baselines across RewardPrediction benchmark. StateFactory establishes a new SOTA among zero-shot methods and approaches supervised baselines. Bold and underline denote the best and second-best results.

Figure 5: StateFactory ablation experiments. We report the EPIC distance ( $D_{\text{EPIC}}$ ) across settings (a-d) and Triplet-based Accuracy for embedding models (e). For metrics, lower $D_{\text{EPIC}}$ and higher Accuracy indicate better performance. Specifically, Triplet-based Accuracy reflects the model’s capability to enforce smaller distances for positive pairs compared to negative ones. Red hatched bars highlight the best performing configurations. Conversely, blue hatched bars indicate the worst performing ones.

Domain	Method	Success Rate (%)
AlfWorld	ReAct ( Yao et al., 2022b )	34.33
AlfWorld	ReAct + StateFactory	55.97
BlocksWorld	ReAct ( Yao et al., 2022b )	85.00
BlocksWorld	ReAct + StateFactory	93.00
ScienceWorld	ReAct ( Yao et al., 2022b )	22.63
ScienceWorld	ReAct + StateFactory	35.03

Table 2: Success Rate on AlfWorld, BlocksWorld, and ScienceWorld. Results show that ReAct + StateFactory consistently improves performance in following domains.

Appendix figures & tables31 assets

Supplementary material from the paper’s appendix.

Appendix

Domain	Pairs (Pos / Neg)	Native Reward Type	Task Description
AlfWorld	336 / 336	Sparse (Binary)	Embodied household instruction following
ScienceWorld	191 / 191	Dense (0-100 Score)	Interactive scientific experiments
WebShop	300 / 300	Dense (Attribute Match)	E-commerce navigation (Subsampled human demos)
BlocksWorld	250 / 250	Sparse (Binary)	Classical planning with spatial reasoning
TextWorld	150 / 150	Sparse (Binary)	Text-adventure games via native engine

Appendix

Table 3: Statistics of the RewardPrediction Benchmark. We report the total number of positive and negative trajectory pairs, the type of native reward signal provided by the environment, and the nature of tasks across different domains.

Step	Action ( $a_{t}$ )	Observation ( $o_{t}$ )	GT ( $R_{t}$ )
0	go to countertop 3	You arrive at countertop 3. On the countertop 3, you see a bread 1… a potato 1…	0.17
1	take potato 1…	You pick up the potato 1 from the countertop 3.	0.33
2	go to microwave 1	You arrive at microwave 1. The microwave 1 is closed.	0.50
3	heat potato 1…	You heat the potato 1 using the microwave 1.	0.67
4	go to diningtable 1	You arrive at diningtable 1. On the diningtable 1, you see a bowl 3…	0.83
5	move potato 1…	You move the potato 1 to the diningtable 1.	1.00

Appendix

Table 4: Visualized Positive Sample (AlfWorld). Task: "put a hot potato in diningtable." The expert segment (

t=0\text{-}5

) receives interpolated dense rewards, while the final random augmentation (

t=6

) inherits the native success signal.

Step	Action ( $a_{t}$ )	Observation ( $o_{t}$ )	GT ( $R_{t}$ )
0	go to stoveburner 1	You arrive at stoveburner 1. On the stoveburner 1, you see a pot 1.	0.00
1	go to cabinet 7	You arrive at cabinet 7. The cabinet 7 is closed.	0.00
2	go to drawer 6	You arrive at drawer 6. The drawer 6 is closed.	0.00
3	open drawer 6	You open the drawer 6. The drawer 6 is open. In it, you see nothing.	0.00
4	go to cabinet 8	You arrive at cabinet 8. The cabinet 8 is closed.	0.00
5	open cabinet 8	You open the cabinet 8. The cabinet 8 is open. In it, you see nothing.	0.00

Appendix

Table 5: Visualized Negative Sample (AlfWorld). Task: "put a hot potato in diningtable." A completely random trajectory where actions fail to advance the task state, resulting in a consistent zero reward.

Parameter	Value
Base Model	Qwen2.5-1.5B
Method	LoRA (SEQ_CLS)
Precision	BF16
Max Length	8192
LoRA Configuration
Rank ( $r$ ) / Alpha ( $\alpha$ )	16 / 32

Appendix

Table 6: Common Hyperparameters. These settings were fixed across all experiments.

Task	Data Config		Compute Setup		Cost
Task	Size	Ep.	GPUs (Mem)	BS	Hours
AlfWorld	35k	3	2 $\times$ A6000 (80G)	16	$\sim$ 4h
TextWorld	35k	3	2 $\times$ A6000 (48G)	4	$\sim$ 4.5h
ScienceWorld	35k	3	2 $\times$ A100 (80G)	16	$\sim$ 6h
BlocksWorld	35k	3	2 $\times$ A100 (80G)	32	$\sim$ 7h
WebShop	35k	2	2 $\times$ A100 (80G)	16	$\sim$ 10h

Appendix

Table 7: Task-Specific Training Details. Comparison of dataset scale, compute resources, and training duration across tasks. Global batch size (BS) scales with GPU count.

Panel of Table 8 — Table 8: Detailed Baseline Prompt Templates. We strictly follow the provided prompt structures for both variants under the Full-History Protocol . The CoT variant (Part I) uses a standard helper system prompt, while the Zero-Shot variant (Part II) explicitly injects a [Without Thinking] constraint in the system prompt to suppress reasoning.

Panel of Table 9 — Table 9: Semantic State Factorization Templates. This pipeline constitutes the first half of the StateFactory method. Module A grounds raw observations into structured states, while Module B filters and refines these states to maintain a coherent, task-relevant history.

Panel of Table 10 — Table 10: Dynamic Goal Interpretation Template. As the second stage of the StateFactory method, this component maintains a persistent yet evolving definition of success, strictly distinguishing between the agent’s “plans” and the task’s “goals”.

Panel of Table 11 — Table 11: Prompt Variations for representation Structural Granularity. We utilized the same backbone and input context across all levels, varying only the Output Schema and Description Constraints to isolate the effect of structural depth.

Panel of Table 12 — Table 12: Polarity Decision Boundaries. We evaluate whether the embedding models correctly rank the semantic paraphrase higher than the distractor. While most models robustly handle spatial contradictions (Case 1), failures in Topic Distraction (Case 2) reveal that some architectures prioritize lexical co-occurrence over factual truth conditions.

State Representation	EPIC Distance ( $\downarrow$ )
State Representation	AlfWorld	ScienceWorld	WebShop	BlocksWorld	TextWorld	Overall
Unstructured Observations (Raw)	0.49	0.45	0.52	0.88	0.50	0.57
Textual States (Flat Denoising)	0.31	0.35	0.42	0.72	0.37	0.43
Object-Centric States (Nested)	0.33	0.33	0.40	0.45	0.26	0.35
Object-Attribute States (Ours)	0.28	0.29	0.29	0.43	0.20	0.30

Appendix

Table 13: Ablation Study I: Representation. We compare varying state representation granularity across five domains. The results demonstrate that structured representations significantly outperform unstructured text, with our Object-Attribute approach achieving the lowest EPIC Distance (

\downarrow

Goal Source Setting	EPIC Distance ( $\downarrow$ )
Goal Source Setting	AlfWorld	ScienceWorld	WebShop	BlocksWorld	TextWorld	Overall
Offline (Oracle Goal State)	0.26	0.27	0.26	0.40	0.22	0.28
Online (Dynamic Goal State)	0.28	0.29	0.29	0.43	0.20	0.30

Appendix

Table 14: Ablation Study II: Goal Interpretation. Comparison between Offline (Oracle) settings using ground-truth goal states and Online (Dynamic) setting using inferred ones. The minimal gap indicates our Dynamic Goal Interpretation is robust without requiring privileged information.

Embedding Backbone	Probing Metric	Downstream Performance (EPIC $\downarrow$ )
Embedding Backbone	Triplet Acc. ( $\uparrow$ )	AlfWorld	ScienceWorld	WebShop	Blocks	TextWorld	Overall
nli-distilroberta-base-v2	0.83	0.28	0.29	0.29	0.47	0.20	0.31
all-mpnet-base-v2	0.76	0.28	0.29	0.29	0.54	0.20	0.32
all-MiniLM-L6-v2	0.81	0.28	0.29	0.29	0.43	0.20	0.30
bge-base-en-v1.5	0.78	0.32	0.34	0.31	0.41	0.25	0.33
embedding-gemma-300m	0.67	0.45	0.47	0.42	0.52	0.38	0.45

Appendix

Table 15: Ablation Study III: Semantic Embeddings. We evaluate StateFactory using five different embedding backbones. We report the EPIC Distance (

\downarrow

) on downstream tasks and the Triplet Accuracy (

\uparrow

) on our synthetic probing dataset. Results show a strong correlation: models with higher Triplet Accuracy (better semantic discriminability) consistently yield lower EPIC distances.

Backbone Model	EPIC Distance ( $\downarrow$ )
Backbone Model	AlfWorld	ScienceWorld	WebShop	Blocks	TextWorld	Overall
Qwen Family (Impact of Scale & Thinking)
Qwen3-14B	0.53	0.57	0.59	0.68	0.53	0.58
Qwen3-30B-A3B-Instruct-2507	0.53	0.59	0.58	0.53	0.52	0.55
Qwen3-30B-A3B-Thinking-2507	0.36	0.46	0.40	0.52	0.31	0.41
GPT-OSS Family (Impact of Reasoning Effort)

Appendix

Table 16: Ablation Study IV: LLM Backbone. We compare different LLM backbones driving the StateFactory pipeline. Lower EPIC distance (

\downarrow

) indicates better alignment with ground-truth progress. The results highlight that “Thinking” variants (Reasoning-Enhanced) and larger parameter scales consistently improve reward precision.

Step	Action	Simplified Obs & State Extraction ( $s_{t}$ )	Goal Interpretation ( $g_{t}$ )	Reward / Gain
1	go to desk 2	Obs: See cd 1 on desk 2. State: {cd 1: [on desk 2]}	Entity Discovery: CD $\to$ cd 1	+0.29 $\uparrow$
…	Steps 2–5: Initial Exploration of the Environment
6	go to desk 1	Obs: See desklamp 1 on desk 1. State: {desklamp 1: [on desk 1]}	Constraint Identification: Lamp $\to$ lamp 1	+0.02 $\uparrow$
7	use desklamp 1	Obs: Lamp 1 is turned ON. State: {lamp 1: [ is turned on ]}	Condition Met: cd 1 under desklamp 1	+0.22 $\uparrow$
8	go to desk 2	Obs: Navigate back to retrieve cd 1 with satisfied environment constraints.	(Unchanged)	+0.31 $\uparrow$
…	Steps 9–11: Redundant Exploration & State Verification

Appendix

Table 17: Trajectory Trace of Case Study 1: “Task: Look at CD under the desklamp.” This case demonstrates the internal logic of the StateFactory engine, specifically how it facilitates dynamic goal interpretation by anchoring abstract instructions to crystallized physical entities (

g_{t}

). By extracting structured states (

s_{t}

) from observations, the engine generates continuous reward gradients (Relative Gain) that guide the MCTS search. Note the significant reward spikes upon environmental constraint satisfaction (Step 7) and the identification of the target entity (Step 8), which effectively steers the agent away from redundant exploration towards a successful execution path.

Step	Action ( $a_{t}$ )	Observation ( $o_{t}$ )	GT ( $R_{t}$ )
0	look around	This room is called the foundry. In it, you see: the agent… a blast furnace…	0.00
1	open door to outside	The door is now open.	0.08
2	go to outside	You move to the outside.	0.17
3	look around	This outside location is called the outside. Here you see: … a baby wolf…	0.17
4	focus on baby baby wolf	You focus on the baby wolf.	0.67
5	pick up baby baby wolf	You move the wolf to the inventory.	0.75

Appendix

Table 18: Visualized Positive Sample (ScienceWorld). Task: "find a(n) living thing… move it to the yellow box." The agent successfully navigates to the outside, identifies the target (baby wolf), and transports it to the goal location.

Step	Action ( $a_{t}$ )	Observation ( $o_{t}$ )	GT ( $R_{t}$ )
0	look at door	ERROR: SHOULD USE OVERRIDE FOR PORTAL.	0.00
1	connect orange to outside	terminal 1 on orange is now connected to terminal 1 on outside	0.00
2	connect air to foundry	terminal 1 on air is now connected to terminal 1 on foundry	0.00
3	connect air to steel table	terminal 2 on air is now connected to terminal 1 on table	0.00
4	open door	The door is now open.	0.08
5	disconnect orange	orange has been disconnected	0.08

Appendix

Table 19: Visualized Negative Sample (ScienceWorld). Task: "find a(n) living thing… move it to the yellow box." The agent performs irrelevant interactions (connecting air, mixing door) that fail to advance the specific task objective.

Step	Action ( $a_{t}$ )	Observation ( $o_{t}$ )	GT ( $R_{t}$ )
0	go to dresser 1	You arrive at dresser 1. On the dresser 1, you see a bowl 1… a cd 2…	0.20
1	take cd 2 from dresser 1	You pick up the cd 2 from the dresser 1.	0.40
2	go to safe 1	You arrive at safe 1. The safe 1 is closed.	0.60
3	open safe 1	You open the safe 1. The safe 1 is open. In it, you see nothing.	0.80
4	move cd 2 to safe 1	You move the cd 2 to the safe 1.	1.00

Appendix

Table 20: Visualized Positive Sample (AlfWorld). Task: "put some cd on safe." The agent identifies the target object (CD 2) on the dresser, navigates to the safe, opens it, and deposits the item.

Step	Action ( $a_{t}$ )	Observation ( $o_{t}$ )	GT ( $R_{t}$ )
0	go to laundryhamper 1	You arrive at laundryhamper 1. On the laundryhamper 1, you see nothing.	0.00
1	go to bed 1	You arrive at bed 1. On the bed 1, you see a cellphone 2…	0.00
2	go to drawer 11	You arrive at drawer 11. The drawer 11 is closed.	0.00
3	look	You are facing the drawer 11. Next to it, you see nothing.	0.00
4	go to drawer 6	You arrive at drawer 6. The drawer 6 is closed.	0.00

Appendix

Table 21: Visualized Negative Sample (AlfWorld). Task: "put some cd on safe." The agent wanders to irrelevant locations (laundry hamper, bed, drawers) and fails to interact with the target objects.

Step	Action ( $a_{t}$ )	Observation ( $o_{t}$ )	GT ( $R_{t}$ )
0	unstack the orange block…	…hand is currently holding orange block, the red block is on the table…	0.25
1	put down the orange block	…hand is empty, the red block is on the table… orange block is on the table…	0.50
2	pick up the red block	…hand is currently holding red block, the blue block is on the table…	0.75
3	stack the red block…	…hand is empty, the red block is on top of the orange block…	1.00

Appendix

Table 22: Visualized Positive Sample (BlocksWorld). Task: "the red block is on top of the orange block." The agent unstack the obstructing orange block to clear the red block, then stacks red onto orange. Rewards are linearly interpolated.

Step	Action ( $a_{t}$ )	Observation ( $o_{t}$ )	GT ( $R_{t}$ )
0	pick up the blue block	…hand is currently holding blue block, the orange block is on top of the red block…	0.00
1	put down the blue block	…hand is empty, the orange block is on top of the red block…	0.00
2	pick up the blue block	…hand is currently holding blue block, the orange block is on top of the red block…	0.00
3	put down the blue block	…hand is empty, the orange block is on top of the red block…	0.00

Appendix

Table 23: Visualized Negative Sample (BlocksWorld). Task: "the red block is on top of the orange block." The agent engages in a repetitive loop with an irrelevant blue block, failing to change the state of the target blocks.

Step	Action ( $a_{t}$ )	Observation ( $o_{t}$ )	GT ( $R_{t}$ )
0	unstack red from orange	…hand is holding red block, the orange block is on top of the blue block…	0.13
1	put down red	…hand is empty, orange block is on top of blue block, red is on table…	0.25
2	unstack orange from blue	…hand is holding orange block, the blue block is on the table…	0.38
3	put down orange	…hand is empty… blue block is on table… orange block is on table…	0.50
4	pick up blue	…hand is holding blue block, the orange block is on the table…	0.63
5	stack blue on orange	…hand is empty, the blue block is on top of the orange block…	0.75

Appendix

Table 24: Visualized Positive Sample (BlocksWorld). Task: "red on blue, blue on orange." The agent dismantles the initial wrong pile (Red on Orange on Blue) and reconstructs it in the correct order.

Step	Action ( $a_{t}$ )	Observation ( $o_{t}$ )	GT ( $R_{t}$ )
0	unstack red from orange	…hand is holding red block, the orange block is on top of the blue block…	0.00
1	stack red on orange	…hand is empty, the red block is on top of the orange block…	0.00
2	unstack red from orange	…hand is holding red block, the orange block is on top of the blue block…	0.00
3	stack red on orange	…hand is empty, the red block is on top of the orange block…	0.00
4	unstack red from orange	…hand is holding red block, the orange block is on top of the blue block…	0.00

Appendix

Table 25: Visualized Negative Sample (BlocksWorld). Task: "red on blue, blue on orange." The agent falls into a reversible loop, repeatedly unstacking and restacking the top block without changing the base structure.

Step	Action ( $a_{t}$ )	Observation ( $o_{t}$ )	GT ( $R_{t}$ )
0	search[gluten free…]	…Banza Chickpea Mac and Cheese… $35.3 … Banza Chickpea Pasta Shells…$ 24.99 …	0.20
1	click[item - banza…]	…Banza Chickpea Pasta Shells… Price: $24.99 … [button] Description … [button] Buy Now	0.40
2	click[description]	…high protein, lower carb, gluten free alternative to traditional mac…	0.60
3	click[< prev]	…Banza Chickpea Pasta Shells… Price: $24.99 … [button] Buy Now	0.80
4	click[buy now]	You have successfully purchased the product.	1.00

Appendix

Table 26: Visualized Positive Sample (WebShop). Task: "i want a gluten free packaged meal, and price lower than 50.00 dollars." The agent correctly filters search results, verifies the item description, and purchases a product ($24.99) meeting the price constraint.

Step	Action ( $a_{t}$ )	Observation ( $o_{t}$ )	GT ( $R_{t}$ )
0	search[i want a gluten…]	…Amy’s Vegan… Burrito… $100.0 … Annie’s Shell Pasta…$ 39.48 …	0.00
1	click[item - B01FW…]	…Amy’s Vegan, Gluten Free Beans & Rice Burrito… Price: $100.0 …	0.00
2	click[Features]	One 5.5 oz. Amy’s Bean & Rice… Gluten Free…	0.00
3	click[Back to Search]	[button] Search [button_]	0.00
4	click[Search]	[button] Search [button_]	0.00

Appendix

Table 27: Visualized Negative Sample (WebShop). Task: "i want a gluten free packaged meal, and price lower than 50.00 dollars." The agent selects an item priced at

100.0 (violating the <

50 constraint) and fails to complete a purchase.

Step	Action ( $a_{t}$ )	Observation ( $o_{t}$ )	GT ( $R_{t}$ )
0	search[sandal ankle…]	…Wedge Sandals… $8.13 … Gibobby Summer Sandals…$ 14.71 …	0.14
1	click[next >]	[button] Back to Search … Page 2 (Total results: 50) …	0.29
2	click[next >]	[button] Back to Search … Page 3 (Total results: 50) …	0.43
3	click[next >]	[button] Back to Search … Page 4 … Gibobby Summer Sandals… $1.87 …	0.57
4	click[item - gibobby…]	…Gibobby Summer Sandals… Price: $1.87 to$ 3.37 … [button] Buy Now	0.71
5	click[black]	…color [clicked button] black… size…	0.86

Appendix

Table 28: Visualized Positive Sample (WebShop). Task: "women ankle strap sandal, black, < 80.00 dollars." The agent navigates through multiple search pages, identifies a compliant item (Gibobby Sandals,

1.87-

3.37), selects the required ’black’ attribute, and completes the purchase.

Step	Action ( $a_{t}$ )	Observation ( $o_{t}$ )	GT ( $R_{t}$ )
0	search[plant based]	…Plant-Based Bolognese… $62.99 … Plant-Based Jerky…$ 4.98 …	0.00
1	click[item - B085…]	…Gardein Ultimate Plant-Based Jerky… Price: $6.29 …	0.00
2	click[teriyaki]	…flavor name… [clicked button] teriyaki…	0.00
3	click[original]	…flavor name… [clicked button] original…	0.00
4	click[Features]	…Contains one 2.25-ounce package of Gardein…	0.00
5	click[< Prev]	…[clicked button] original… [clicked button] teriyaki…	0.00

Appendix

Table 29: Visualized Negative Sample (WebShop). Task: "women ankle strap sandal, black, < 80.00 dollars." The agent performs a completely irrelevant search ("plant based") and interacts with food items (Jerky) instead of footwear, failing the task constraints entirely.

Step	Action ( $a_{t}$ )	Observation ( $o_{t}$ )	GT ( $R_{t}$ )
0	examine type T box	The type T box looks strong, and impossible to destroy…	0.00
1	take cuboid keycard	You pick up the cuboid keycard from the ground.	0.20
2	unlock cuboid chest…	You unlock the cuboid chest.	0.40
3	open cuboid chest	You open the cuboid chest, revealing a potato and a knife.	0.60
4	take knife…	You take the knife from the cuboid chest.	0.80
5	put knife on board	You put the knife on the board… The End	1.00

Appendix

Table 30: Visualized Positive Sample (TextWorld). Task: "lift the cuboid keycard… take the knife… place the knife on the board." After an initial inspection, the agent executes a perfect sequence: getting the key, unlocking the chest, retrieving the knife, and placing it on the target.

Step	Action ( $a_{t}$ )	Observation ( $o_{t}$ )	GT ( $R_{t}$ )
0	inventory	You are carrying: a mop.	0.00
1	put mop on board	You put the mop on the board.	0.00
2	close fridge	You close the fridge.	0.00
3	examine mop	The mop would seem to be out of place here	0.00
4	take type T keycard	You pick up the type T keycard from the ground.	0.00
5	examine board	The board is stable. On the board is a mop.	0.00

Appendix

Table 31: Visualized Negative Sample (TextWorld). Task: "lift the cuboid keycard… take the knife… place the knife on the board." The agent interacts with irrelevant objects (mop, fridge) and picks up the wrong keycard (Type T instead of Cuboid), failing to progress.

Step	Action ( $a_{t}$ )	Observation ( $o_{t}$ )	GT ( $R_{t}$ )
0	take key	You pick up the key from the ground.	0.20
1	unlock safe with key	(the safe with the key) You unlock the safe.	0.40
2	open safe	(the safe) You open the safe, revealing a spoon.	0.60
3	take spoon from safe	(the safe) You take the spoon from the safe.	0.80
4	insert spoon…	You put the spoon into the spherical safe… The End	1.00

Appendix

Table 32: Visualized Positive Sample (TextWorld). Task: "retrieve key… unlock safe… insert spoon." The agent successfully identifies the key, unlocks the container, and retrieves the target tool to complete the puzzle.

Step	Action ( $a_{t}$ )	Observation ( $o_{t}$ )	GT ( $R_{t}$ )
0	examine saucepan	The saucepan is stable.	0.00
1	examine coffee cup	The coffee cup is cheap looking.	0.00
2	examine spherical latchkey	The spherical latchkey is cold to the touch	0.00
3	eat legume	You eat the legume. Not bad.	0.00
4	examine saucepan	The saucepan is stable.	0.00

Appendix

Table 33: Visualized Negative Sample (TextWorld). Task: "retrieve key… unlock safe… insert spoon." The agent gets distracted by irrelevant objects (saucepan, coffee cup) and consumes food instead of solving the lock puzzle.

Reward Modeling

Reward modeling learns scores or rankings that capture human goals and task success, guiding AI selection, reinforcement learning, and alignment.

See all 325 papers

Oct 8, 2026

CausalDreamer: Learning Predictive World Models with Latent Disentanglement

World models for control must capture which aspects of the environment respond to the agent's actions and which are relevant to reward. Generative world models such as Dreamer 4 consist of a video tokenizer, which encodes each frame into a latent, and a dynamics model, which is pretrained to predict future latents from past latents and actions. Yet the tokenizer is trained with a reconstruction objective, without action or reward supervision, so its latent provides no explicit mechanism to separate controllable, uncontrollable, reward-relevant, and reward-irrelevant information. We propose CausalDreamer, which keeps the tokenizer frozen and re-encodes its latent into a factored representation of four groups along two axes: controllability, where only the two controllable groups receive the action, and reward relevance, learned by predicting the reward from the two reward-relevant groups. The pretrained dynamics model is then fine-tuned to predict the factored representation. We evaluate CausalDreamer and the pretrained world model it starts from with model-predictive planning on 20 MMBench2 tasks: 10 clean tasks seen during training and 10 unseen tasks, of which 6 are manipulated variants of clean tasks with a changed background, object, or maze layout, and 4 are new environments. We normalize returns so that a policy taking uniformly random actions scores 0 and an expert scores 1. CausalDreamer achieves a 14% higher normalized score than the pretrained world model on the clean tasks (0.199 vs.\ 0.175) and a 25% higher score on the manipulated variants (0.307 vs.\ 0.246), while neither model scores meaningfully above the random policy in the new environments. Additionally, our analysis shows that the factored representation separates reward-irrelevant changes, such as a changed background, from its reward-relevant groups.

Prince Jha, Nils Lukas, Kun Zhang +1

Mohamed bin Zayed University of Artificial Intelligence (MBZUAI) · Carnegie Mellon University

World Model Learning Causal World Models Disentangled Representation Learning Model-Based Planning Causal Representation Learning

Ask PDF

Jul 30, 2026

World Action Planner: Generalizable Decision-Making with Action-Conditioned World Models

Building generalizable agents for diverse applications remains a fundamental challenge. While imitation learning-based policies succeed in specific training environments, they often fail to generalize to novel scenes and tasks. In this work, we propose World Action Planner, a robot planning system that leverages the reasoning capabilities of Vision-Language Models (VLMs) and the physical grounding of a multi-task pose-image conditioned world model. Our system enables an agent to propose initial action plans and iteratively refine them via optimization and search, reasoning over imagined world model rollouts. We demonstrate that our approach achieves superior performance across compositional tasks, new layouts, and zero-shot generalization scenarios, significantly outperforming state-of-the-art end-to-end policy models such as VLAs and WAMs. Project website at worldactionplanner.github.io

Xiangcheng Zhang, Yilun Du

Harvard University

World Models for Robotics World Model-Based Planning Action-Conditioned World Models Robot Policy Generalization World Models for VLA Models

Ask PDF

Sep 21, 2026

World State Generator

Language agents solve complex tasks through plans and actions. A single step the world refuses puts the goal out of reach, and what the agent does next decides the task. Prompted planners fail at exactly this point, rewriting the refused step in new words, meeting the same refusal, and burning the attempt budget without moving. They fail because the plan was never tied to the world, so a refusal has nothing in the plan to attach to. A world is where a task runs, and it has its own rules, its own admissible actions, and its own constraints. We build synthetic worlds across 7 domains and extract training data from them. A program enforces each world's rules and grades its goal, and every world is admitted only if its goal is reachable from its initial state. Agents run inside and leave verified failures paired with repairs that carried the run to a state the world certified, a record of about 226K trajectories. On this record we train the World State Generator, a model that writes a plan as checkable states of the world and keeps that plan aligned with the world it runs in. That alignment is what a plan written in language lacks, since the world it runs in has physical limits, logical dependencies, and required orders the language never states, and the plan encounters these rules only when a state fails. WSG takes that failure as the rule the world has stated and rewrites the remaining states to obey it, so the plan bends to the world as the run goes on. Across 7 public benchmarks, WSG raises end-to-end success for two open models near 30B parameters over prompting and brings to the level of proprietary model.

Sungheon Jeong, Sanggeon Yun, Ryozo Masukawa +3

University of California, Irvine · Northeastern University

LLM Agent Planning LLM Planning World Model-Based Planning

Ask PDF

Abstract

Reward Modeling

Explore similar work

CausalDreamer: Learning Predictive World Models with Latent Disentanglement

World Action Planner: Generalizable Decision-Making with Action-Conditioned World Models

World State Generator