Mental World Modeling

# Mental World Modeling: Technical Breakdown

Core Idea§

The paper introduces Mental World Modeling (MWM), a framework that extends traditional world models by explicitly modeling hidden mental states (beliefs, desires, intentions) alongside physical states. This captures how agents perceive, reason, and act in social scenarios. The key insight is that predicting human behavior requires tracking what each agent knows and believes about the world, not just the physical scene.

Formalization§

Let the world state at time $t$ be a tuple $(\mathcal{P}_t, \mathcal{M}_t)$ where $\mathcal{P}_t$ is the physical state (positions, objects, etc.) and $\mathcal{M}_t$ is the set of mental states for each agent $i$: $\mathcal{M}_t = \{m^i_t\}$ with $m^i_t = (\text{beliefs}^i_t, \text{desires}^i_t, \text{intentions}^i_t, \text{feelings}^i_t)$. The world model $W$ predicts the next state:

$$(\mathcal{P}_{t+1}, \mathcal{M}_{t+1}) = W(\mathcal{P}_t, \mathcal{M}_t, \mathbf{a}_t)$$

where $\mathbf{a}_t$ is the joint action of all agents. To render observations, MWM uses a target-specific partial observation function $O$ that determines what a particular agent perceives:

$$o^i_t = O(\mathcal{P}_t, \mathcal{M}_t, i)$$

This enables simulating how actions affect both physical reality and each agent's mental model of the world.

MENTIS Baseline§

MENTIS is a training-free instantiation of MWM that decomposes inference into five modules:

1. State Parsing: Convert raw story (text/image/video) into initial $\mathcal{P}_0$ and $\mathcal{M}_0$ using an LLM-based parser. 2. Target-Observation Generation: For each agent $i$, generate $o^i_t = \text{LLM}(\mathcal{P}_t, \mathcal{M}_t, i)$. 3. Action Decomposition: Enumerate candidate actions $\mathbf{a}_t$ using an LLM that proposes plausible actions based on agents' mental states. 4. Coupled Physical-Mental Transition: For each candidate action, predict the next state: $$(\mathcal{P}_{t+1}, \mathcal{M}_{t+1}) = \text{LLM}_{\text{trans}}(\mathcal{P}_t, \mathcal{M}_t, \mathbf{a}_t)$$ 5. Branch-Level Value Evaluation: Score each trajectory based on how well it aligns with observed human decisions (e.g., using a value function $V$).

Code Illustration§

Below is a simplified PyTorch-like pseudocode illustrating the coupled transition step (assuming an LLM as a callable oracle):

import torch
from typing import Dict, Any

def coupled_transition(physical_state: Dict, mental_states: Dict, action: Dict, llm: Callable) -> tuple:
    # Serialize state and action into a prompt
    prompt = f"Physical: {physical_state}\nMental: {mental_states}\nAction: {action}\nNext state:"
    response = llm.generate(prompt)
    # Parse response (assume LLM returns structured output)
    next_physical, next_mental = parse_state(response)
    return next_physical, next_mental

# Example usage
physical = {'agent_positions': {'Alice': (0,0), 'Bob': (1,1)}, 'objects': {'ball': (2,2)}}
mental = {'Alice': {'beliefs': {'ball_at': (2,2)}, 'desires': 'get_ball'},
          'Bob': {'beliefs': {'ball_at': (2,2)}, 'desires': 'pass_ball'}}
action = {'Alice': 'move_to(2,2)', 'Bob': 'kick_ball_toward(Alice)'}
next_phys, next_ment = coupled_transition(physical, mental, action, gpt4)

Evaluation§

The authors construct a dataset of situated decision scenarios (text, image, video) and evaluate 8 LLM-based world models. They compare models with and without explicit mental state modeling. Results show that MWM significantly improves prediction accuracy, with ablation studies revealing bottlenecks in state parsing and transition fidelity. The framework is training-free and fully inspectable, making it suitable for debugging and interpretation.

Key Contributions§

• Formalizes mental states as first-class citizens in world models.
• Proposes a training-free baseline MENTIS that can be implemented with off-the-shelf LLMs.
• Demonstrates the necessity of mental modeling for social decision prediction.
• Identifies current limitations (e.g., LLMs struggle with recursive belief modeling).