Causal world models under physical parameter shifts

Chu Lin; Yi Sun

cs.LGOct 21, 2026

Causal world models under physical parameter shifts

Authors: Chu Lin, Yi Sun

Organizations: Department of Computer Science, The City College of New York, New York, NY 10031, USA · Department of Electrical Engineering, The City College of New York, New York, NY 10031, USA

Abstract

World models predict how embodied systems evolve under actions, but their accuracy can degrade when physical conditions move outside the training distribution. We study two components of causal modeling: which inputs a predictor may use, and how its mechanism computes their effects. In a synthetic benchmark with shortcut features, a frozen dependency mask derived from paired interventions reduces mean out-of-distribution (OOD) RMSE from 1.769 for a fully connected predictor to 0.697, close to the generator-oracle mask (0.694). For robot link-mass extrapolation in CausalWorld, with disjoint training and test mass ranges, we develop finger-local world models in which link masses enter inertia and dynamics-bias terms through explicit affine coefficients, actions act through either a learned torque network or the known PD controller, and learned residuals capture contact-related effects. Compared with a separately trained control whose mass inputs are fixed at training means, the known-PD model turns out-of-range masses into a consistent predictive gain: median time-averaged joint R2 over 150-step rollouts rises from 0.609 to 0.727, with improvement in all ten task-seed runs (exact two-sided sign test, p = 0.002). In a shared 100-step forecast, the learned-torque model is competitive with the strongest recurrent robotic world model (RWM) configuration (median joint R2 of 0.755 - 0.768 versus 0.752), whereas supplying out-of-range masses sharply degrades the MLP and recurrent predictors. These results indicate that how a physical parameter enters the dynamics mechanism determines whether its out-of-range value helps or harms prediction.

Figures & tables

Figure 1: Overview of the robot world model. Three finger mechanisms predict joint accelerations under fixed input dependencies; one mechanism family is used per model. A separate object MLP predicts object-state changes. Joint integration and object updates produce future states, which are fed back for rollout prediction. Training uses simulator trajectories; evaluation uses link masses outside the training range.

Setting	One-step training	MLP / learned torque rollout	Known-PD rollout
Optimizer / batch	Adam / 256	Adam / 128	Adam / 128
Learning rate	$10^{-3}$	$3\times 10^{-4}$	$\{10^{-4},3\times 10^{-5}\}$
Horizon	1	50	10, then 50
Training budget	80 epochs	1,500 updates	1,500 updates per stage
Validation interval / patience	Every epoch / 10	Every 100 updates / 5	Every 25 updates / 8

Table 1: Robot training settings. Actual-mass and fixed-mass-input PD models use the same candidate grid and select their settings separately on light-mass validation data.

Input mask	OOD RMSE
Dense	1.769
Random sparse	1.604
Intervention-derived, frozen	0.697
Generator oracle reference	0.694

Table 2: Synthetic OOD prediction. Values are ten-seed means in training-standardized output coordinates. Lower is better.

Mechanism	Actual mass	Ablated mass	Ablation definition
MLP	$-0.803$	0.721	Mass inputs removed
Mass-affine, learned torque	0.734	0.733	Mass inputs fixed at training means
Mass-affine, known PD	0.727	0.609	Mass inputs fixed at training means

Table 3: Mass-information ablation under link-mass extrapolation. Values are medians of 150-step mean joint

R^{2}

over two tasks and five seeds. Within each row, the actual-mass model is compared with a separately trained control of the same mechanism family. Higher is better.

Panel of Figure 2 — Figure 2: Predictive value of actual mass within the known-PD mechanism. (a) Points are medians of per-step scores across the ten task–seed runs; lines connect recorded evaluation steps. (b) Bars show medians of within-run differences in 150-step mean joint $R^{2}$ . These paired differences are aggregated separately from the model-score medians in Table 3 .

Model	Mass input	$\overline{R^{2}}\uparrow$	Joint	Object RMSE	$\Delta_{\mathrm{mass}}\uparrow$
			NRMSE $\downarrow$	(cm) $\downarrow$
Hold-state reference	None	$-0.320$	1.055	3.835	—
RWM ( $N_{\mathrm{tr}}=8$ )	None	0.718	0.478	2.913	—
	Actual	$-0.648$	1.182	7.311	$\downarrow\,1.367$
RWM ( $N_{\mathrm{tr}}=50$ )	None	0.752	0.457	2.842	—
	Actual	$-0.540$	1.136	5.744	$\downarrow\,1.293$

Table 4: CausalWorld link-mass extrapolation: a shared 100-step forecast after 32 observed history steps. Entries are medians over ten task–seed runs.

N_{\mathrm{tr}}

denotes RWM’s training forecast length. The final column is the median paired gain from actual mass relative to the control immediately above it. Bold marks the best value in each metric column.

Predictor	Pushing	Picking	Overall
PETS-style probabilistic ensemble	0.2321	0.0492	0.1407
PETS-style ensemble + graph input selection	0.6831	0.3839	0.5335
Intervention-masked predictor	0.7508	0.4737	0.6122

Table 5: External dynamics-baseline comparison on CausalWorld space A

\rightarrow

B. Values are mean prediction skill relative to the source-training mean,

\mathcal{S}

(higher is better), over ten seeds per task. All rows share the same data, targets, and scoring procedure. This is a complementary environment-transfer benchmark.

Appendix figures & tables2 assets

Supplementary material from the paper’s appendix.

Appendix

Component	Established foundation	Work developed in this study
Input dependencies	Causal dynamics and dependency selection [ 4 , 10 ]	Frozen intervention masks in the synthetic benchmark; shared finger-local constraints in robot models.
Dynamics mechanisms	Positive-definite inertia parameterization and inertial identification [ 15 , 17 , 18 ]	Mass-affine inertia and state-dependent bias; torque and contact-related residual networks receive no direct mass input.
Known control	Hybrid modeling and simulator PD control [ 19 , 18 , 26 , 27 ]	A known-PD mechanism in physical units, integrated with learned mass-affine terms and an object-state predictor.
Evaluation	Controlled simulation and input ablation	Disjoint link-mass ranges, separately trained mass-information controls, and comparison of ID and OOD mass benefits.

Appendix

Table 6: Sources of established components and this study’s specific design and evaluation choices.

Model	Mass input	ID	OOD pushing	OOD picking
Hold-state reference	None	-0.332	-0.327	-0.316
RWM (8)	None	0.642	0.715	0.719
RWM (8)	Actual	0.626	-0.850	-0.447
RWM (50)	None	0.635	0.752	0.751
RWM (50)	Actual	0.622	-0.658	-0.421
Mass-affine, learned torque (A)	Mean	0.701	0.764	0.774

Appendix

Table 7: Additional shared-window scores. ID columns aggregate ten task–seed runs; each OOD task column aggregates its five seeds. All scores are median 100-step mean joint

R^{2}

(higher is better).

Causal World Models

Causal world models represent how entities and environments change under actions, enabling prediction, planning, and counterfactual reasoning.

See all 46 papers

Feedback World Model Enables Precise Guidance of Diffusion Policy
May 15, 2026Tuo An, Jindou Jia, Gen Li +8World Model Learning World Models for Robotics
MARS Lab, Nanyang Technological University
PhysiCalWorld: In-Context Physics Calibration for World Models
Oct 21, 2026Yunhao Luo, Wentao He, Yinpei Dai +5World Models Action-Conditioned World Models
University of Michigan
Do-JEPA: From Masking to Intervention in Latent World Models
Sep 29, 2026Hossein Resani, Javen Qinfeng ShiCausal Effect Estimation Causal World Models
Australian Institute for Machine Learning, Adelaide University · Responsible AI Research Centre

May 15, 2026J/K move · Enter open · S save

Feedback World Model Enables Precise Guidance of Diffusion Policy

Tuo An, Jindou Jia, Gen Li, Jingliang Li+7

MARS Lab, Nanyang Technological University

World Model Learning World Models for Robotics Diffusion Policy Robot Policy Generalization Inference-Time Guidance

Open paper AskPDF arXiv

World models aim to improve robotic decision making by predicting the consequences of actions. However, in practice, their predictions often become unreliable once the robot encounters states outside the training distribution, limiting their effectiveness at deployment. We observe that execution itself provides a natural but underutilized signal: after each action, the robot directly observes the true next state, revealing the mismatch between predicted and actual outcomes. Building on this insight, we propose feedback world model, an inference-time mechanism that closes the loop between prediction and observation at inference time. Instead of treating the world model as a static open-loop predictor, our method maintains a lightweight feedback state that is updated online to iteratively correct future predictions, compensating for model errors using real-time observations without additional training data or parameter updates. We show that this process can be interpreted as a latent-space observer and admits convergence guarantees under bounded, asymptotically convergent dynamics residuals and a stable feedback gain. We further introduce action-aware guidance to better translate corrected predictions into control by emphasizing action-controllable components while suppressing irrelevant variations. Experiments on LIBERO-Plus, Robomimic, and real-world manipulation tasks demonstrate that our method substantially improves both prediction accuracy and policy performance under distribution shift, reducing world model prediction error by up to 76.4%. The same feedback mechanism improves the predictions of the pretrained foundation model Fast-WAM-Joint without updating its parameters, demonstrating its plug-and-play applicability. These results show that incorporating real-time feedback at inference time provides a simple yet powerful alternative to static world modeling.

Abstract

Causal World Models

Explore similar work

Feedback World Model Enables Precise Guidance of Diffusion Policy

PhysiCalWorld: In-Context Physics Calibration for World Models

Do-JEPA: From Masking to Intervention in Latent World Models