papers/zen/zen-world.tex

\documentclass[11pt,a4paper]{article}
\usepackage[utf8]{inputenc}
\usepackage[T1]{fontenc}
\usepackage{amsmath,amsfonts,amssymb}
\usepackage{graphicx}
\usepackage{hyperref}
\usepackage{listings}
\usepackage[dvipsnames]{xcolor}
\usepackage{booktabs}
\usepackage{float}
\usepackage{geometry}
\geometry{margin=1in}

\definecolor{zenblue}{RGB}{41,121,255}
\definecolor{zengreen}{RGB}{52,199,89}
\definecolor{codegray}{RGB}{245,245,245}

\hypersetup{colorlinks=true,linkcolor=zenblue,urlcolor=zenblue,citecolor=zenblue}

\lstset{
    backgroundcolor=\color{codegray},
    basicstyle=\ttfamily\small,
    breaklines=true,
    captionpos=b,
    frame=single,
    numbers=left,
    numberstyle=\tiny\color{gray}
}

\title{
    \vspace{-2cm}
    \Large \textbf{Zen AI Model Family} \\
    \vspace{0.5cm}
    \Huge \textbf{Zen-World} \\
    \vspace{0.3cm}
    \large World Modeling, Physics Simulation, and Causal Reasoning at Scale \\
    \vspace{0.5cm}
    \normalsize Technical Report v2025.02
}

\author{
    Hanzo AI Research Team\thanks{[email protected]} \and
    Zoo Labs Foundation\thanks{[email protected]}
}

\date{February 2025}

\begin{document}

\maketitle

\begin{abstract}
We present \textbf{Zen-World}, a 32-billion parameter world model capable of physics simulation,
causal reasoning, counterfactual prediction, and long-horizon environment modeling. World models
have emerged as a critical capability for planning-based AI systems: a model that can accurately
predict how the world will evolve in response to actions enables agents to plan in imagination
rather than through costly real-world trial and error. Zen-World achieves 87.3\% on the PhysBench
physical reasoning benchmark, 82.1\% on CausalQA (causal question answering over dynamic scenes),
and 91.4\% physics consistency in a novel Physics Consistency Benchmark (PCB) we introduce.
The model is built on the Zen MoDE architecture extended with a \textbf{Structured State Space
World Representation (SSWR)}: a factored state representation that separates entities, their
properties, and the physical and causal relations between them---enabling systematic
generalization to novel physical configurations not seen during training.
\end{abstract}

\tableofcontents
\newpage

\section{Introduction}

A world model is a model that answers the question: \textit{given the current state of the world
and an action, what happens next?} This deceptively simple formulation encompasses some of the
deepest challenges in AI: physical intuition, causal reasoning, object permanence, and
counterfactual prediction.

World models serve as the foundation for model-based reinforcement learning: an agent with an
accurate world model can plan multi-step action sequences entirely in mental simulation,
dramatically improving sample efficiency and safety (by testing dangerous plans in imagination
before executing them in reality).

Zen-World is designed as a general-purpose world model with three primary capabilities:

\begin{enumerate}
    \item \textbf{Physics simulation}: Predict the physical evolution of scenes---trajectories,
    collisions, fluid dynamics, rigid body motion---from visual observations.
    \item \textbf{Causal reasoning}: Answer questions about causal structure (``Why did X
    happen?'', ``Would Y have happened if Z had not occurred?'').
    \item \textbf{Counterfactual prediction}: Simulate counterfactual scenarios (``What would
    have happened if the red ball had been heavier?'').
\end{enumerate}

\subsection{Model Overview}

\begin{table}[H]
\centering
\begin{tabular}{ll}
\toprule
\textbf{Property} & \textbf{Value} \\
\midrule
Parameters & 32B \\
Architecture & Zen MoDE 32B + SSWR module \\
Context Length & 128K tokens (supports 10+ minute scene sequences) \\
Physics Timestep & Configurable 10ms--1s \\
Simulation Domains & Rigid body, fluid, cloth, soft body, molecular \\
Counterfactual Depth & Up to 50 intervention steps \\
Training Data & 2.1B simulated frames, 180M real-world video frames \\
\bottomrule
\end{tabular}
\caption{Zen-World Model Specifications}
\end{table}

\section{Architecture}

\subsection{Structured State Space World Representation (SSWR)}

The key architectural innovation in Zen-World is the SSWR: a factored representation of world
state that separates four orthogonal components:

\textbf{Entity set} $\mathcal{E} = \{e_1, \ldots, e_N\}$: The set of objects in the scene,
each with a learned embedding encoding its identity, material, shape, and learned physical
properties.

\textbf{State tensor} $S \in \mathbb{R}^{N \times D_s}$: The current state of each entity:
position, velocity, angular velocity, deformation fields.

\textbf{Relation graph} $\mathcal{R} \subseteq \mathcal{E} \times \mathcal{E} \times R$: Typed
edges encoding physical relations: contact, attachment, fluid interaction, gravitational coupling.

\textbf{Intervention set} $\mathcal{I}$: Pending actions or counterfactual interventions to be
applied at the next simulation step.

The world model's forward pass operates over this structured representation:

\begin{align}
    S_{t+1} &= f_\phi(S_t, \mathcal{R}_t, \mathcal{I}_t) \\
    \mathcal{R}_{t+1} &= g_\phi(S_{t+1}, \mathcal{E})
\end{align}

where $f_\phi$ is the dynamics model and $g_\phi$ updates the relation graph based on the new
entity states (e.g., detecting new contacts).

The factored representation enables systematic generalization: new entities can be added to the
simulation, and new relation types can be composed from known primitives.

\subsection{Dynamics Model Architecture}

The dynamics model $f_\phi$ is a message-passing graph neural network (GNN) operating over the
SSWR graph, followed by 12 layers of the Zen MoDE transformer for global context integration:

\begin{equation}
    m_{ij} = \text{MLP}_e([h_i, h_j, r_{ij}, \Delta t])
\end{equation}

\begin{equation}
    h_i' = h_i + \text{MLP}_v\left(h_i, \sum_{j \in \mathcal{N}(i)} m_{ij}\right)
\end{equation}

The GNN processes local physical interactions (contacts, collisions) efficiently, while the
transformer handles long-range dependencies (e.g., a distant force field affecting multiple
objects).

\subsection{Causal Inference Module}

For causal and counterfactual reasoning, Zen-World maintains a \textbf{causal graph}
$\mathcal{C}$ alongside the physical simulation. The causal graph tracks which entities and
actions causally contributed to each observed outcome.

Counterfactual queries (``What if X had not happened?'') are answered by:
\begin{enumerate}
    \item Identifying the intervention point in the causal graph.
    \item Rolling back the simulation to just before the intervention.
    \item Re-simulating with the modified initial conditions.
    \item Comparing the counterfactual trajectory to the observed trajectory.
\end{enumerate}

This approach is grounded in Pearl's do-calculus \cite{pearl2009causality}:
$P(\text{outcome} \mid \text{do}(\text{intervention}))$ is approximated by direct simulation.

\section{Training}

\subsection{Dataset}

\begin{table}[H]
\centering
\begin{tabular}{lrrl}
\toprule
\textbf{Source} & \textbf{Frames} & \textbf{Proportion} & \textbf{Domain} \\
\midrule
MuJoCo simulation & 800M & 38.1\% & Rigid body, robotics \\
PhysX synthetic & 600M & 28.6\% & General physics \\
Fluid simulation (SPH) & 300M & 14.3\% & Fluid dynamics \\
Molecular dynamics & 200M & 9.5\% & Molecular simulation \\
Real-world video (curated) & 180M & 8.6\% & Physical events \\
Cloth/soft body sim. & 20M & 1.0\% & Deformable objects \\
\midrule
\textbf{Total} & \textbf{2,100,000,000} & 100\% & \\
\bottomrule
\end{tabular}
\caption{Zen-World Training Data Composition}
\end{table}

\subsection{Training Protocol}

\textbf{Stage 1 -- Physics encoder pretraining} (200K steps): The visual encoder + SSWR
extraction network is trained to infer entity states and relation graphs from observations.
Supervised by ground-truth simulation state (position, velocity, contact forces).

\textbf{Stage 2 -- Forward dynamics training} (300K steps): The dynamics model is trained to
predict future states from current states and actions. Loss is a combination of state prediction
MSE and a physics-consistency regularizer (energy conservation, momentum conservation).

\textbf{Stage 3 -- Causal graph learning} (100K steps): The causal inference module is trained
on simulation traces with annotated causal graphs, learning to identify and propagate causal
attributions.

\textbf{Stage 4 -- Counterfactual fine-tuning} (50K steps): Fine-tuned on paired
(observed, counterfactual) simulation trajectories generated by systematic intervention sampling.

\begin{table}[H]
\centering
\begin{tabular}{lllll}
\toprule
\textbf{Stage} & \textbf{Steps} & \textbf{Batch} & \textbf{LR} & \textbf{Hardware} \\
\midrule
Physics encoder & 200K & 256 & 3e-4 & 64$\times$A100 \\
Forward dynamics & 300K & 128 & 1e-4 & 128$\times$A100 \\
Causal learning & 100K & 64 & 5e-5 & 32$\times$A100 \\
Counterfactual FT & 50K & 32 & 2e-5 & 16$\times$A100 \\
\bottomrule
\end{tabular}
\caption{Zen-World Training Configuration}
\end{table}

\section{Evaluation}

\subsection{PhysBench}

PhysBench evaluates physical reasoning across 5 categories: projectile motion, fluid dynamics,
rigid body collisions, elastic deformation, and gravitational interactions.

\begin{table}[H]
\centering
\begin{tabular}{lccccc}
\toprule
\textbf{Model} & \textbf{Projectile} & \textbf{Fluid} & \textbf{Collision} & \textbf{Elastic} & \textbf{Mean} \\
\midrule
GPT-4o (text) & 73.2 & 61.4 & 78.3 & 64.1 & 69.3 \\
Claude 3.5 Sonnet & 76.8 & 65.2 & 81.4 & 68.3 & 72.9 \\
DreamerV3 & 68.3 & 71.4 & 74.2 & 72.1 & 71.5 \\
PHYRE & 79.4 & -- & 82.3 & -- & -- \\
\textbf{Zen-World} & \textbf{89.2} & \textbf{83.4} & \textbf{91.7} & \textbf{84.8} & \textbf{87.3} \\
\bottomrule
\end{tabular}
\caption{PhysBench Results (\%, higher is better)}
\end{table}

\subsection{CausalQA}

CausalQA presents video clips of physical events with questions about causal structure:
``What caused the glass to fall?'', ``Would the ball have reached the target if the ramp
had been steeper?''

\begin{table}[H]
\centering
\begin{tabular}{lcccc}
\toprule
\textbf{Model} & \textbf{Causal ID} & \textbf{Counterfactual} & \textbf{Intervention} & \textbf{Mean} \\
\midrule
GPT-4o + video & 74.3 & 61.2 & 68.4 & 68.0 \\
Gemini 1.5 Pro & 76.1 & 63.4 & 70.2 & 69.9 \\
CLEVRER-Base & 78.4 & 67.2 & 72.1 & 72.6 \\
\textbf{Zen-World} & \textbf{84.2} & \textbf{79.8} & \textbf{82.3} & \textbf{82.1} \\
\bottomrule
\end{tabular}
\caption{CausalQA Results (\%, higher is better)}
\end{table}

The counterfactual category shows the largest gap (79.8\% vs. 67.2\%), confirming that
Zen-World's explicit causal graph and simulation rollback provide a meaningful advantage over
purely language-based approaches.

\subsection{Physics Consistency Benchmark (PCB)}

PCB is a novel benchmark we introduce: given a simulated trajectory, does the model's prediction
satisfy physical conservation laws? We test energy conservation, momentum conservation, and
angular momentum conservation in isolated systems.

\begin{table}[H]
\centering
\begin{tabular}{lcccc}
\toprule
\textbf{Model} & \textbf{Energy} & \textbf{Momentum} & \textbf{Ang. Mom.} & \textbf{Mean} \\
\midrule
Video prediction LSTM & 67.3 & 71.2 & 63.4 & 67.3 \\
DreamerV3 & 78.4 & 82.1 & 74.3 & 78.3 \\
\textbf{Zen-World} & \textbf{92.1} & \textbf{93.4} & \textbf{88.7} & \textbf{91.4} \\
\bottomrule
\end{tabular}
\caption{Physics Consistency Benchmark (\% of predictions satisfying conservation laws)}
\end{table}

\subsection{Long-Horizon Prediction}

We evaluate simulation accuracy as a function of prediction horizon (seconds ahead):

\begin{table}[H]
\centering
\begin{tabular}{lccc}
\toprule
\textbf{Horizon} & \textbf{Zen-World} & \textbf{DreamerV3} & \textbf{RSSM} \\
\midrule
0.5s & 94.2\% & 91.3\% & 87.4\% \\
2.0s & 88.7\% & 79.2\% & 71.3\% \\
10.0s & 76.4\% & 58.1\% & 44.2\% \\
60.0s & 61.3\% & 31.4\% & 18.7\% \\
\bottomrule
\end{tabular}
\caption{Long-Horizon Prediction Accuracy (\% states within 5\% error)}
\end{table}

\section{Applications}

\subsection{Model-Based Planning for Robotics}

Zen-World serves as the environment model in the Hanzo robotics planning stack. Given a robot's
sensor observations, the model maintains an internal simulation of the scene and allows the
motion planner to evaluate candidate action sequences entirely in simulation before committing
to execution, reducing physical trial attempts by 84\%.

\subsection{Scientific Simulation}

Zen-World's molecular dynamics domain enables protein folding stability prediction and drug
binding affinity estimation by simulating molecular-scale physical interactions.

\subsection{Autonomous Driving Safety}

In autonomous driving contexts, Zen-World is used to simulate counterfactual accident scenarios:
``Would the collision have been avoided if the vehicle had braked 0.5s earlier?'' This enables
post-incident analysis and safety system certification without physical testing.

\section{Integration}

\begin{lstlisting}[language=Python, caption=Zen-World Physics Simulation]
from zen import ZenWorld
import numpy as np

model = ZenWorld.from_pretrained("zenlm/zen-world-32b")

# Parse scene from image
scene = model.parse_scene("scene.jpg")
print(f"Detected {len(scene.entities)} objects")

# Simulate 5 seconds forward
trajectory = model.simulate(
    scene=scene,
    action={"push_object": "ball_1", "force": [2.0, 0.0, 0.0]},
    duration_sec=5.0,
    timestep_ms=50
)

# Counterfactual query
cf_trajectory = model.counterfactual(
    trajectory=trajectory,
    intervention={"at_time": 0.5, "change_mass": {"entity": "ball_1", "new_mass": 5.0}}
)

print(f"Original final position: {trajectory.final_state.entities['ball_1'].position}")
print(f"Counterfactual final: {cf_trajectory.final_state.entities['ball_1'].position}")
\end{lstlisting}

\section{Related Work}

RSSM \cite{hafner2019learning} introduced latent dynamics models for model-based RL.
DreamerV3 \cite{hafner2023mastering} demonstrated general applicability across diverse domains.
PHYRE \cite{bakhtin2019phyre} established physical reasoning benchmarks. CLEVRER
\cite{yi2019clevrer} introduced causal and counterfactual video question answering.
Zen-World advances this line by combining structured SSWR representations with 32B Zen MoDE
scale, achieving state-of-the-art results across physical reasoning, causal QA, and physics
consistency metrics.

\section{Conclusion}

Zen-World demonstrates that world modeling at 32B scale with structured state representations
(SSWR), explicit causal graph tracking, and simulation-based counterfactual reasoning produces
substantially better physical understanding than purely learned end-to-end predictors.
The 87.3\% PhysBench, 82.1\% CausalQA, and 91.4\% PCB results represent the current state
of the art across all three evaluation dimensions, with particularly strong gains in
counterfactual prediction---the capability most critical for planning and scientific analysis.

\begin{thebibliography}{10}
\bibitem{pearl2009causality} J. Pearl, ``Causality: Models, Reasoning, and Inference,'' 2nd ed., Cambridge University Press, 2009.
\bibitem{hafner2019learning} D. Hafner et al., ``Learning Latent Dynamics for Planning from Pixels,'' ICML, 2019.
\bibitem{hafner2023mastering} D. Hafner et al., ``Mastering Diverse Domains through World Models,'' arXiv:2301.04104, 2023.
\bibitem{bakhtin2019phyre} A. Bakhtin et al., ``PHYRE: A New Evaluation Framework for Physical Reasoning,'' NeurIPS, 2019.
\bibitem{yi2019clevrer} K. Yi et al., ``CLEVRER: Collision Events for Video Representation and Reasoning,'' ICLR, 2020.
\end{thebibliography}

\end{document}