Causal world models under physical parameter shifts
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
| Setting | One-step training | MLP / learned torque rollout | Known-PD rollout |
|---|---|---|---|
| Optimizer / batch | Adam / 256 | Adam / 128 | Adam / 128 |
| Learning rate | |||
| 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 |
| Input mask | OOD RMSE |
|---|---|
| Dense | 1.769 |
| Random sparse | 1.604 |
| Intervention-derived, frozen | 0.697 |
| Generator oracle reference | 0.694 |
| Mechanism | Actual mass | Ablated mass | Ablation definition |
|---|---|---|---|
| MLP | 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 |
| Model | Mass input | Joint | Object RMSE | ||
| NRMSE | (cm) | ||||
| Hold-state reference | None | 1.055 | 3.835 | — | |
| RWM ( ) | None | 0.718 | 0.478 | 2.913 | — |
| Actual | 1.182 | 7.311 | |||
| RWM ( ) | None | 0.752 | 0.457 | 2.842 | — |
| Actual | 1.136 | 5.744 |
| 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 |
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. |
| 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 |
Causal World Models
Causal world models represent how entities and environments change under actions, enabling prediction, planning, and counterfactual reasoning.
See all 46 papers