Reciprocity as an Ante-Hoc Inductive Bias for Physics-Consistent World Models

Abstract

Learned world models are an efficient alternative to numerical simulators, but their rollouts are hard to audit against known physical laws. The XAI4Science workshop asks for architectures whose explanations can be checked in this way, not only for weather and climate foundation models but for scientific dynamics more broadly. We study a deliberately small setting: a two-dimensional symplectic-Euler plant whose configuration is a nonzero complex number $z$, together with a residual multilayer perceptron that predicts the next state given the current state and a thrust. A single conservation identity—for every existing state $x\ne 0$, $x\cdot (1/x)=1$—is encoded ante hoc as input features $|z|$ and $1/|z|$, as a residual product penalty, as a hinge that forbids the unrepresentable singularity at 0, and as a soft energy match against a restoring well at $|z|=1$. On a recorded seed-7 run (numpy MLP 8→32→32→4, trained in 1.54 s), post-training on embodied tasks reduces held-out prediction MSE from 0.00243 to 0.00163, open-loop horizon collisions from 0.10 to 0.00, and closed-loop reach success from 0.833 to 0.917, while the axiom violation rate remains 0. The contribution is an evaluable conservation prior and a physics-first evaluation protocol, not a weather foundation model and not a claim of scale.

1. Introduction

Foundation models for weather and climate now match operational numerical forecasts at a fraction of the compute of classical simulators [Lam et al., 2023, Pathak et al., 2022]. The same pattern appears in control: compact world models support planning and policy learning without querying an expensive plant at every imagination step [Ha and Schmidhuber, 2018, Hafner et al., 2020, 2025, LeCun, 2022]. Efficiency, however, is bought with opacity. Unlike a physics-based integrator, whose state update is an instance of a known equation, a learned rollout can look plausible while quietly violating conservation, entering a forbidden set, or hallucinating a trajectory that no downstream user should trust [Montavon et al., 2018]. XAI4Science puts this gap at the centre of the programme. The workshop asks for three complementary threads: ante-hoc interpretability and inductive biases; post-hoc attribution with a faithfulness check; and physics-consistent explanations that can be validated against known laws.

This paper takes the first and third threads seriously in a setting small enough that every number can be recomputed. We do not train a weather foundation model, and we do not claim state-of-the-art control. We encode one algebraic identity as a hard modelling choice,

and then ask whether a learned action-conditioned world model can be checked against that identity while remaining useful for control. The identity is not a decorative regularizer. In the plant, 0 is unrepresentable: inversion of the origin raises an error, and any trajectory segment that would enter a disk of radius 0.07 is refused and restored by a radial map that preserves $\arg (z)$ and sends $|z|$ to 1. The unit circle $|z|=1$ is a restoring homeostasis well, not an arbitrary task target. Time evolution is a symplectic Euler step [Hairer et al., 2006] on a potential that combines that well with a $1/|z{|}^2$ barrier.

A residual MLP with 1,476 parameters is pretrained on free dynamics and then post-trained on four embodied tasks (reach a pose on the unit circle, survive a near-horizon approach, return home after a shove, invert $z\mapsto 1/z$ and restabilize). Control uses short-horizon shooting in the learned model; the ground-truth plant still takes the physical step. Evaluation follows the workshop rather than pixel fidelity: we report the algebraic gap of (1), the rate at which open-loop rollouts hit the forbidden horizon, energy consistency, and closed-loop success on the plant.

Contributions. (i) An ante-hoc recipe that turns a single conservation identity into features, a loss, and a forbidden set for an action-conditioned world model. (ii) A physics-first evaluation protocol that separates physical correctness, safety, and downstream control utility. (iii) A complete recorded experiment (seed 7) showing that task post-training improves prediction and safety without introducing axiom violations, with every metric taken from that run and no additional trials invented for this paper.

2. Related work

World models for control

Ha and Schmidhuber [2018] showed that a compact learned dynamics model can stand in for the environment during policy evolution. Dreamer plans in a latent world model [Hafner et al., 2020]; DreamerV3 scales the same idea across many control domains with a single configuration [Hafner et al., 2025]. Joint-embedding predictive architectures likewise treat a world model as the substrate of autonomous intelligence [LeCun, 2022]. These lines are typically judged on task return or reconstruction. They rarely expose a conservation identity that a scientist can evaluate independently of reward.

Structure-preserving learned dynamics

Physics-informed networks penalize PDE residuals [Raissi et al., 2019]. Hamiltonian and Lagrangian neural networks hard-wire energy structure into the parameterization [Greydanus et al., 2019, Cranmer et al., 2020]. Neural ODEs learn a vector field that is then integrated [Chen et al., 2018]; symplectic ODE-Net targets Hamiltonian systems with control [Zhong et al., 2020]. Sparse identification recovers governing equations from data [Brunton et al., 2016, Rudy et al., 2017]. We use the same scientific instinct—prefer models that cannot easily lie about a known law—but attach it to a residual world model used for model-predictive control, and we score the model on law-gap, forbidden-set entry, and closed-loop utility together.

Foundation models and scientific XAI

GraphCast and FourCastNet demonstrate that learned weather models can be operationally competitive [Lam et al., 2023, Pathak et al., 2022]. The XAI4Science call notes that this efficiency introduces a transparency problem: physical hallucinations are difficult to trace. Post-hoc attribution is one response [Montavon et al., 2018]; ante-hoc inductive biases are another. Our testbed is not a climate model. It is a controlled analogue in which the law, the singularity, and the restoring well are fully known, so that a conservation prior can be evaluated without confounding by data assimilation or subgrid parameterization.

3. A reciprocity prior for world models

3.1 Plant

The body is a unit-mass point whose configuration is $z=x+iy\in \mathbb{C}\setminus \{0\}$. The state is $s=(x,y,v_x,v_y)$ and the action is a clipped thrust $a=(a_x,a_y)$ with $\Vert a\Vert \le 2.4$. The potential is

so $|z|=1$ is a radial restoring well (homeostasis) and the second term is a barrier against the origin. A linear drag of coefficient 0.16 is included in the force, so the plant is near-symplectic rather than exactly Hamiltonian: symplectic Euler [Hairer et al., 2006] with step $\Delta t=0.05$ updates velocity then position,

If the line segment $[z,z^{+}]$ meets the disk $|z|<0.07$ (the horizon), or if $z^{+}$ is non-finite, the step is a cataclysm: it is refused, $|z|$ is renormalized to 1 at constant argument, and a fraction of the velocity is reversed. A concentric danger disk of radius 0.22 is used only as a cost and as a near-miss task region; it is not a hard constraint. The origin itself cannot be stored. The discrete invert primitive $z\mapsto 1/z$ maps velocity as $v\mapsto -v/z^2$ and is itself refused if the segment would cross the horizon. After every legal step the product $z\cdot (1/z)$ equals 1 up to floating-point dust (on the order of $10^{-16}$). That is the ground-truth statement of (1) in this plant.

3.2 World model

The world model is a two-hidden-layer tanh MLP of width 32, implemented in numpy, with a residual head:

The last two coordinates are the reciprocity features: they make the scale $\mu =|z|$ and its inverse first-class inputs, so the network does not have to rediscover that $1/|z|$ diverges as the body approaches the forbidden set. Inputs are standardized from the training mixture. The architecture is 8→32→32→4 (1,476 parameters). This is an ante-hoc bias in the sense of the workshop: the identity is in the interface of the model, not recovered afterwards by an explainer.

3.3 Loss

For a batch of transitions $(s,a,s^{\prime })$ the training objective is

with $({\lambda }_a,{\lambda }_h,{\lambda }_e)=(0.25,0.80,0.04)$. ${\mathcal{L}}_{\mathrm{MSE}}$ is a weighted squared error that emphasizes position over velocity (weights 1.4 and 0.85). The axiom term ${\mathcal{L}}_{\mathrm{ax}}$ is the mean squared gap $|\hat{z}\cdot (1/\hat{z})-1{|}^2$, with a large penalty if $|\hat{z}|$ is below machine epsilon. The hinge ${\mathcal{L}}_h$ penalizes $\max (0,0.07+0.04-|\hat{z}|{)}^2$, i.e. predicted states that approach or enter the horizon. The energy term ${\mathcal{L}}_E$ is the mean squared mismatch between $T+V$ on ${\hat{s}}_{t+1}$ and on the true $s^{\prime }$, using (2). Adam (batch 64) is used at learning rate $3\times 10^{-3}$ in pretraining and $1.1\times 10^{-3}$ in post-training, with gradient clipping at 5.

A comment is required for honesty. In exact arithmetic, $z\cdot (1/z)=1$ for every finite $z\ne 0$, so ${\mathcal{L}}_{\mathrm{ax}}$ is numerically at roundoff whenever the prediction stays away from the origin. The term is still the right interface to the law: it screams when a prediction becomes illegal or non-finite, and it is the quantity we later report as a physics check. The hinge and the energy match carry the non-vacuous geometric content (stay outside the horizon; match $T+V$).

3.4 Two-stage training and control

Pretraining uses 3,200 transitions of random thrusts on bodies spawned with $|z|\in [0.48,3.1]$, i.e. free coverage of the plane away from 0 (420 optimizer steps). Post-training copies those weights and continues for 560 steps on 3,600 transitions from four embodied tasks, mixed with occasional pretrain minibatches to limit forgetting: (i) reach a random unit-circle pose from $|z|\in [1.45,2.7]$; (ii) near-miss, starting inside the danger disk with inward velocity and driving out; (iii) homeostasis after a random shove from the unit circle; (iv) invert then restabilize. A held-out set of 900 Physical-AI transitions is used for prediction metrics and is never trained on.

Closed-loop control does not clone an expert PD controller. At each plant step the agent scores a fixed palette of candidate thrusts (coast, two gains toward the target, radial escape and approach, two tangents, and eight compass directions) by rolling the world model for $H=8$ steps and summing distance-to-target, a speed penalty, a danger hinge, a $|z|-1$ penalty, and a large cost if a predicted state is non-finite or inside the horizon. The minimizing first action is applied to the ground-truth plant. Figure 1 summarizes the loop; Figure 2 shows the geometry the loop is about.

4. Evaluation protocol

Metrics are chosen to match the workshop’s physics-consistency thread rather than reconstruction of pixels or latent log-likelihoods.

Physical correctness. On the 900 held-out transitions we report one-step prediction MSE (and a position/velocity split in the appendix), the axiom violation rate (a prediction is a violation if it is non-finite, has $|\hat{z}|<0.07$, or has product-gap $>10^{-6}$), the mean product-gap, and the mean absolute energy error between ${\hat{s}}_{t+1}$ and $s^{\prime }$.

Safety. Open-loop horizon collisions: 40 near-miss initial conditions, each rolled $H=8$ steps in the world model under a restoring PD action, counting the fraction of rollouts that predict $|\hat{z}|$ below the horizon. This is not closed-loop plant safety; it asks whether the model imagines entering the forbidden set.

Control utility. Shooting MPC ($H=8$) on the plant, with fixed evaluation seed. Reach: 12 episodes, success if the body comes within 0.18 of a unit-circle target while remaining outside the danger disk (horizon 90 steps). Home: 8 episodes, success if $||z|-1|<0.12$ and speed $<0.85$ after a shove. Near-miss: 8 episodes, success if the body never hits the horizon and finishes at $|z|>0.35$. Invert-home: 8 episodes, success if after $z\mapsto 1/z$ the body returns to $||z|-1|<0.16$. Embodied success is the unweighted mean of the four task rates (36 episodes in total).

All numbers below are from one recorded run (seed 7, 1.54 s wall time, no extra fine-tune rounds). We do not average over seeds in this submission, and we do not invent further experiments.

5. Experiments

Table 1 and Figure 3 report the recorded metrics. Post-training cuts held-out one-step MSE from 0.0024258668165594243 to 0.0016316616372299964 (a relative drop of about one third) and energy MAE from 0.03212 to 0.02478. The axiom violation rate is 0 at both stages; the mean product-gap stays at floating-point roundoff (5.33e-17 pretrained, 5.79e-17 post-trained). Open-loop horizon collisions fall from 0.10 (4/40) to 0.00 (0/40). On the plant, reach rises from 0.833 (10/12) to 0.917 (11/12); home, near-miss, and invert-home are already 8/8 after pretraining and remain so; embodied success therefore moves from 0.944 (34/36) to 0.972 (35/36). No plant cataclysm occurred in either closed-loop evaluation. Mean best reach distance improved from 0.199 to 0.165 (appendix).

Three readings follow, in decreasing order of how much the sample size will bear.

First, the conservation identity is not the quantity that post-training “improves,” and we do not claim that it is. Both models already satisfy (1) at roundoff on held-out one-step predictions. What post-training changes is the geometry of the predictor: one-step error and energy match get tighter, and the model stops imagining $H=8$ trajectories that fall into the forbidden disk. That is the physics-consistency result that the hinge and energy terms are for.

Second, the control improvement is real in the recorded run and small in absolute terms. Reach is one additional success on twelve episodes; embodied success is one additional success on thirty-six. A reader who wants a confidence interval will find that 10/12 and 11/12 overlap substantially. We therefore treat the closed-loop rates as a descriptive utility check—the world model is usable for shooting MPC—not as a claim that post-training is a statistically dominant control method. The more reliable scalar in the table is the n=900 prediction MSE.

Third, pretraining on free dynamics is already enough for home, near-miss, and invert-home. The restoring well in (2) does a large fraction of the work, which is precisely what an ante-hoc physical bias is supposed to do: some tasks become easy because the plant, and then the model, cannot wander far from the law. Reach, which requires moving to a specific argument on the unit circle from outside, is the task that still needs post-training.

6. What would transfer

The plant is a toy. The protocol is not. A weather or climate world model cannot honestly claim to “have encoded reciprocity” because of this paper, and we do not make that claim. What does transfer is the evaluation contract that XAI4Science is asking for.

Ante-hoc: pick a law the scientific community already agrees on (dry-air mass, energy, non-negative humidity, a well-posed equation of state) and put it in the architecture or the loss before any explainer is trained, in the same way $|z|$ and $1/|z|$ are first-class inputs here [Raissi et al., 2019, Greydanus et al., 2019]. Physics-consistent scoring: report the residual of that law on held-out rollouts, not only RMSE against reanalysis [Lam et al., 2023, Pathak et al., 2022]. Safety: define the unphysical set explicitly (here, the disk $|z|<0.07$) and count how often imagined trajectories enter it; a climate analogue is a rollout that produces negative water or an unbounded tendency. Utility: keep a downstream task that is not reconstruction—here, shooting MPC on the true plant; in a forecasting FM, a decision or nowcast that a scientist actually uses. Post-hoc attribution [Montavon et al., 2018] remains complementary: once the law-gap is measured, one can still ask which inputs drove a violating rollout. The miniature experiment is a unit test for that contract, not a substitute for GraphCast-scale training.

7. Limitations

This is a controlled two-dimensional symplectic testbed, not a weather or climate foundation model, and not a robotics benchmark. The body is a point mass; there are no contacts, no partial observations, and no pixels. The world model is a tiny MLP with 1,476 parameters, trained in under two seconds on CPU. Closed-loop evaluation uses 12+8+8+8 episodes from one seed. Drag makes the integrator dissipative, so we should not be read as claiming exact energy conservation or exact unitarity—only a near-symplectic discrete step plus an algebraic identity that holds off the origin. The axiom gap, as discussed in Section 3, is a roundoff-level check for illegal predictions rather than a sensitive residual of a PDE. We did not ablate the four loss terms in a factorial design; the appendix only reports the split of MSE into position and velocity already stored in the recorded metrics file. The contribution we are willing to defend is the evaluable conservation prior and the evaluation protocol, not scale, not SOTA, and not a drop-in replacement for GraphCast-class models.

8. Conclusion

XAI4Science asks for scientific models whose claims can be checked against known laws. We encoded one such law—reciprocity $x\cdot (1/x)=1$ for $x\ne 0$—as an ante-hoc inductive bias on a learned world model, with 0 as a forbidden singularity and $|z|=1$ as a restoring well. On a fully recorded miniature experiment, post-training improved one-step prediction and eliminated open-loop horizon collisions without ever producing an axiom violation, and a shooting controller that imagined in the learned model remained competent on the true plant. The same protocol—report the law-gap, the forbidden-set rate, and the downstream utility, and refuse to hide behind reconstruction error—is available for larger scientific world models. We offer this testbed as a small, inspectable instance of that protocol, not as a surrogate for a foundation model.

References

1. Steven L. Brunton, Joshua L. Proctor, and J. Nathan Kutz. Discovering governing equations from data by sparse identification of nonlinear dynamical systems. Proceedings of the National Academy of Sciences, 113(15):3932–3937, 2016.
2. Ricky T. Q. Chen, Yulia Rubanova, Jesse Bettencourt, and David Duvenaud. Neural ordinary differential equations. In Advances in Neural Information Processing Systems, volume 31, 2018.
3. Miles Cranmer, Sam Greydanus, Stephan Hoyer, Peter Battaglia, David Spergel, and Shirley Ho. Lagrangian neural networks. arXiv preprint arXiv:2003.04630, 2020.
4. Samuel Greydanus, Misko Dzamba, and Jason Yosinski. Hamiltonian neural networks. In Advances in Neural Information Processing Systems, volume 32, 2019.
5. David Ha and Jürgen Schmidhuber. Recurrent world models facilitate policy evolution. In Advances in Neural Information Processing Systems, volume 31, 2018.
6. Danijar Hafner, Timothy Lillicrap, Jimmy Ba, and Mohammad Norouzi. Dream to control: Learning behaviors by latent imagination. In International Conference on Learning Representations, 2020.
7. Danijar Hafner, Jurgis Pasukonis, Jimmy Ba, and Timothy Lillicrap. Mastering diverse control tasks through world models. Nature, 640:647–653, 2025.
8. Ernst Hairer, Christian Lubich, and Gerhard Wanner. Geometric Numerical Integration: Structure-Preserving Algorithms for Ordinary Differential Equations. Springer, 2 edition, 2006.
9. Remi Lam, Alvaro Sanchez-Gonzalez, Matthew Willson, Peter Wirnsberger, Meire Fortunato, Ferran Alet, Suman Ravuri, Timo Ewalds, Zach Eaton-Rosen, Weihua Hu, Alexander Merose, Stephan Hoyer, George Holland, Oriol Vinyals, Jacklynn Stott, Alexander Pritzel, Shakir Mohamed, and Peter Battaglia. Learning skillful medium-range global weather forecasting. Science, 382(6677):1416–1421, 2023.
10. Yann LeCun. A path towards autonomous machine intelligence. OpenReview, 2022.
11. Grégoire Montavon, Wojciech Samek, and Klaus-Robert Müller. Methods for interpreting and understanding deep neural networks. Digital Signal Processing, 73:1–15, 2018.
12. Jaideep Pathak, Shashank Subramanian, Peter Harrington, Sanjeev Raja, Ashesh Chattopadhyay, Morteza Mardani, Thorsten Kurth, David Hall, Zongyi Li, Kamyar Azizzadenesheli, Pedram Hassanzadeh, Karthik Kashinath, and Animashree Anandkumar. FourCastNet: A global data-driven high-resolution weather model using adaptive Fourier neural operators. arXiv preprint arXiv:2202.11214, 2022.
13. Maziar Raissi, Paris Perdikaris, and George E. Karniadakis. Physics-informed neural networks: A deep learning framework for solving forward and inverse problems involving nonlinear partial differential equations. Journal of Computational Physics, 378:686–707, 2019.
14. Samuel H. Rudy, Steven L. Brunton, Joshua L. Proctor, and J. Nathan Kutz. Data-driven discovery of partial differential equations. Science Advances, 3(4):e1602614, 2017.
15. Yaofeng Desmond Zhong, Biswadip Dey, and Amit Chakraborty. Symplectic ODE-net: Learning Hamiltonian dynamics with control. In International Conference on Learning Representations, 2020.

A. Additional recorded metrics and training details

The recorded metrics file also stores a position/velocity split of one-step error, a mean best reach distance, and a plant-cataclysm count. We reproduce those numbers here and do not add new trials. Position error is already small after pretraining; most of the one-step MSE is in velocity, and both halves improve after post-training. One-step predictions never enter the horizon at either stage (horizon-pred. rate 0): the 0.10→0.00 collision result is specifically about open-loop multi-step imagination, which is the relevant quantity for shooting MPC.

B. Training curve of the recorded run

Figure 4 redraws the 1-step training MSE that was logged every 20 optimizer steps during the recorded run. Pretraining is noisy, including one large batch spike, and then sits near the held-out level 0.00243. Post-training continues from the copied weights and ends near the held-out level 0.00163. The figure is diagnostic, not an additional experiment.