Learning Continuous Neural Representation of Stochastic Hybrid Systems

Sangli Teng; Hang Liu; Koushil Sreenath

cs.LGSep 30, 2026

Learning Continuous Neural Representation of Stochastic Hybrid Systems

Authors: Sangli Teng, Hang Liu, Koushil Sreenath

Abstract

A stochastic hybrid system (SHS) is governed by a stochastic differential equation (SDE) describing the continuous dynamics and a Markov reset kernel triggered on the guard surface. Its probability evolution can be described by a hybrid Fokker-Planck (HFP) equation with a partial differential term corresponding to the SDE and an integral term arising from the reset kernel. This work shows that such an SHS can be approximated by an SDE in a higher-dimensional latent space where the sample paths are continuous. The key to this result is to encode different branches of the reset kernel using auxiliary variables, transforming the resets into deterministic ones that enable topological gluing. By the embedding theorem, the glued manifold can then be embedded into a higher-dimensional Euclidean space. We show that the probability evolution on the embedded image no longer requires explicit reset terms in the HFP equation. Building on this theorem, we design a loss that matches the evolving state distributions, enabling a single latent SDE to recover the probability evolution of the SHS without mode labeling, trajectory segmentation, or event-based simulations.

Figures & tables

Figure 1: We approximate SHSs with a Markov reset kernel by an A-SHS with a deterministic reset function that allows the topological gluing to make the hybrid flow continuous.

Figure 2: Target distribution. The dashed lines indicate the boundary of the partition of unity.

Figure 3: The density evolution of A-SHS for ( 15 ). We find that the marginal density of A-SHS in

x

matches that of ( 15 ). As shown in the

x-v

plot, the density is mixed by the OU process first and then redistributed to the GMM to approximate the reset by Construction 1 and 2 .

Figure 5: The predicted probability evolution of bouncing ball systems under GMM reset. Our method correctly captures the long-term distribution with a training horizon of

0.32s

Figure 6: We plotted five predicted trajectories of the bouncing ball system under GMM reset conditioned on the same initial condition. The reset function independently chooses restitution by a GMM centered at

0.5

0.9

at each collision. Our method correctly captures this pattern.

	Ours	( Teng et al., 2026b )	( Li et al., 2020 )	( Li et al., 2020 ) $\star$	( Bartosh et al., 2025 )	( Kiyohara et al., 2025 )
b-ball (GMM)	$\mathbf{0.257}\;(\mathbf{0.214})$	Diverged	$0.369\;(1.12)$	$0.338\;(0.879)$	$0.620\;(4.14)$	$0.594\;(4.22)$
b-ball (Uniform)	$\mathbf{0.289}\;(\mathbf{0.249})$	Diverged	$0.404\;(1.03)$	$0.380\;(0.873)$	$0.732\;(4.84)$	$0.542\;(3.23)$
Torus (GMM)	$\mathbf{0.603}\;(\mathbf{0.333})$	Diverged	$1.84\;(6.17)$	$4.32\;(11.3)$	Diverged	$0.909\;(6.27)$
Torus (Uniform)	$\mathbf{0.560}\;(\mathbf{0.304})$	Diverged	$2.67\;(40.5)$	$1.51\;(10.3)$	Diverged	$0.944\;(7.19)$
Klein (GMM)	$\mathbf{0.612}\;(\mathbf{0.341})$	Diverged	Diverged	Diverged	$10.2\;(256)$	$0.903\;(6.05)$
Klein (Uniform)	$\mathbf{0.560}\;(\mathbf{0.298})$	Diverged	$1.33\;(15.1)$	$1.96\;(29.2)$	$7.84\;(193)$	$0.919\;(6.66)$

Table 1: Baseline comparison of conditional (unconditional) path distribution loss across 11 benchmark problems. Per-case means and sample standard deviations are reported in Tables 5 and 6 .

\star

denotes using

\mathcal{L}_{x}

and

\mathcal{L}_{\mathrm{x}}^{\mathrm{traj}}

. Horizon:

1\,\mathrm{s}

Setting	$d_{z}=d_{x}$	$d_{z}=2d_{x}$	$d_{z}=3d_{x}$	MLP	$\mathcal{L}_{\mathrm{x}}^{\mathrm{traj,c}}$	$\mathcal{L}_{z}$	$\mathcal{L}_{x}$	$\mathcal{L}_{\mathrm{x}}^{\mathrm{traj}}$	$\mathcal{L}_{\mathrm{x}}^{\mathrm{traj,c}},\ \mathcal{L}_{\mathrm{x}}^{\mathrm{traj}}$	$\mathcal{L}_{x},\ \mathcal{L}_{\mathrm{x}}^{\mathrm{traj}}$	$\mathcal{L}_{\mathrm{x}}^{\mathrm{traj,c}},\ \mathcal{L}_{x}$
Mean change (%)	$+14.1\;(+71.4)$	$+2.9\;(+12.2)$	$+0.3\;(+4.7)$	$+3.6\;(+14.7)$	$+6.2\;(+6.6)$	$+19.5\;(+246.6)$	$+1.2\;(+8.8)$	$+0.0\;(+2.1)$	$+23.3\;(+54.7)$	$-0.7\;(+3.8)$	$+5.7\;(+4.4)$
Worse cases (/11)	$11\;(11)$	$8\;(6)$	$5\;(5)$	$8\;(8)$	$7\;(7)$	$11\;(11)$	$7\;(9)$	$6\;(5)$	$11\;(11)$	$5\;(7)$	$6\;(5)$

Table 2: Ablation summary across 11 benchmark settings on the conditional (unconditional) trajectory distribution loss. Mean change is the average of per-setting percentage changes relative to the default model; positive values indicate degradation. Worse cases count the settings with a higher mean test loss than the default model (out of 11), using the means over five training seeds. MLP replaces the diffusion sampler with a deterministic MLP encoder; loss labels identify the removed objectives. Detailed results are presented in Tables 11 , 13 , 12 , 14 and 15 .

Appendix figures & tables36 assets

Supplementary material from the paper’s appendix.

Appendix

Configuration	$H$	$L$	$C$	$M$	$F$
Standard	256	128	64	100	64
Large	512	256	128	200	128

Appendix

Table 3: Hidden-width settings. The large configuration is used only for the two planar balls.

H

is the hidden width for Proposed and CHyLL;

L

and

C

are the hidden and context widths for Latent SDE;

M

is the hidden width for SDE Matching;

F

is the hidden width for NSF.

Component	Proposed	CHyLL	Latent SDE	SDE Matching	NSF
Initial-state network	Score: $[H]\times 2$	$[H]\times 2$	$[L]\times 2$	$[M]\times 1$	Conditional flow
Posterior context	N/A	N/A	GRU( $L$ ), context $C$	GRU( $M$ ), $[M]\times 2$	Bridge model
Prior drift	$[H]\times 2$	$[H]\times 2$	$[L]\times 2$	$[M]\times 1$	N/A
Diffusion	$[H]\times 2$ , matrix	N/A	$[L]\times 1$ , per coordinate	$[M]\times 1$ , per coordinate	N/A
Observation decoder	$[H]\times 3$	$[H]\times 3$	Linear	Linear	Observed-state flow
Flow conditioner	N/A	N/A	N/A	N/A	$[F]\times 2$

Appendix

Table 4: Network components used in the baseline comparison.

[h]\times\ell

denotes

\ell

hidden layers of width

h

; output layers are not counted. Widths are specified in Table 3 . Latent SDE with additional losses has the same architecture as Latent SDE.

Case

Ours

( Teng et al., 2026b )

( Li et al., 2020 )

\star

( Bartosh et al., 2025 )

( Kiyohara et al., 2025 )

b-ball (GMM)

\mathbf{0.257}

\,\pm\,

0.0158

Diverged

0.369

\,\pm\,

0.0464

0.338

\,\pm\,

0.0269

0.620

\,\pm\,

0.265

0.594

\,\pm\,

0.00740

b-ball (Uniform)

\mathbf{0.289}

\,\pm\,

0.0317

Diverged

0.404

\,\pm\,

0.0856

0.380

\,\pm\,

0.0984

0.732

\,\pm\,

0.295

0.542

\,\pm\,

0.00676

Torus (GMM)

\mathbf{0.603}

\,\pm\,

0.00347

Diverged

1.84

\,\pm\,

2.48

4.32

\,\pm\,

7.81

Diverged

0.909

\,\pm\,

0.0214

Torus (Uniform)

\mathbf{0.560}

\,\pm\,

0.00756

Diverged

2.67

\,\pm\,

3.93

1.51

\,\pm\,

1.60

Diverged

0.944

\,\pm\,

0.0184

Klein (GMM)

\mathbf{0.612}

\,\pm\,

0.0124

Diverged

10.2

\,\pm\,

8.15

0.903

\,\pm\,

0.0120

Klein (Uniform)

\mathbf{0.560}

\,\pm\,

0.0178

Diverged

1.33

\,\pm\,

0.835

1.96

\,\pm\,

2.27

7.84

\,\pm\,

6.17

0.919

\,\pm\,

0.0175

Appendix

Table 5: Baseline comparison measured by test-set conditional path distribution loss

\mathcal{L}_{\mathrm{x}}^{\mathrm{traj,c}}

(

\downarrow

). Horizon:

1\,\mathrm{s}

Case

Ours

( Teng et al., 2026b )

( Li et al., 2020 )

\star

( Bartosh et al., 2025 )

( Kiyohara et al., 2025 )

b-ball (GMM)

\mathbf{0.214}

\,\pm\,

0.0558

Diverged

1.12

\,\pm\,

0.571

0.879

\,\pm\,

0.293

4.14

\,\pm\,

3.75

4.22

\,\pm\,

0.0888

b-ball (Uniform)

\mathbf{0.249}

\,\pm\,

0.0819

Diverged

1.03

\,\pm\,

0.708

0.873

\,\pm\,

0.673

4.84

\,\pm\,

4.54

3.23

\,\pm\,

0.0527

Torus (GMM)

\mathbf{0.333}

\,\pm\,

0.0133

Diverged

6.17

\,\pm\,

9.07

11.3

\,\pm\,

16.3

Diverged

6.27

\,\pm\,

0.551

Torus (Uniform)

\mathbf{0.304}

\,\pm\,

0.0319

Diverged

40.5

\,\pm\,

76.1

10.3

\,\pm\,

15.1

Diverged

7.19

\,\pm\,

0.492

Klein (GMM)

\mathbf{0.341}

\,\pm\,

0.0201

Diverged

256

\,\pm\,

231

6.05

\,\pm\,

0.330

Klein (Uniform)

\mathbf{0.298}

\,\pm\,

0.0317

Diverged

15.1

\,\pm\,

17.5

29.2

\,\pm\,

49.4

193

\,\pm\,

171

6.66

\,\pm\,

0.451

Appendix

Table 6: Baseline comparison measured by test-set unconditional path distribution loss

\mathcal{L}_{\mathrm{x}}^{\mathrm{traj}}

(

\downarrow

). Horizon:

1\,\mathrm{s}

Case	Ours		( Teng et al., 2026b )	( Li et al., 2020 )		( Li et al., 2020 ) $\star$		( Bartosh et al., 2025 )	( Kiyohara et al., 2025 )
b-ball (GMM)	$\mathbf{0.487}$ $\,\pm\,$	$0.0153$	Diverged	$1.60$ $\,\pm\,$	$0.647$	$1.18$ $\,\pm\,$	$0.390$	Diverged	$0.897$ $\,\pm\,$	$0.0116$
b-ball (Uniform)	$\mathbf{0.433}$ $\,\pm\,$	$0.0152$	Diverged	$1.42$ $\,\pm\,$	$1.04$	$1.45$ $\,\pm\,$	$1.70$	Diverged	$0.693$ $\,\pm\,$	$0.00792$
Torus (GMM)	$\mathbf{0.671}$ $\,\pm\,$	$0.00239$	Diverged	Diverged		Diverged		Diverged	$0.913$ $\,\pm\,$	$0.0227$
Torus (Uniform)	$\mathbf{0.654}$ $\,\pm\,$	$0.00691$	Diverged	Diverged		Diverged		Diverged	$0.954$ $\,\pm\,$	$0.0201$
Klein (GMM)	$\mathbf{0.677}$ $\,\pm\,$	$0.00530$	Diverged	Diverged		Diverged		Diverged	$0.905$ $\,\pm\,$	$0.0150$
Klein (Uniform)	$\mathbf{0.653}$ $\,\pm\,$	$0.0108$	Diverged	Diverged		Diverged		Diverged	$0.922$ $\,\pm\,$	$0.0189$

Appendix

Table 7: Baseline comparison measured by test-set conditional path distribution loss

\mathcal{L}_{\mathrm{x}}^{\mathrm{traj,c}}

(

\downarrow

). Horizon:

3\,\mathrm{s}

Case	Ours		( Teng et al., 2026b )	( Li et al., 2020 )		( Li et al., 2020 ) $\star$		( Bartosh et al., 2025 )	( Kiyohara et al., 2025 )
b-ball (GMM)	$\mathbf{0.732}$ $\,\pm\,$	$0.137$	Diverged	$33.0$ $\,\pm\,$	$21.3$	$19.3$ $\,\pm\,$	$12.8$	Diverged	$9.62$ $\,\pm\,$	$0.281$
b-ball (Uniform)	$\mathbf{0.690}$ $\,\pm\,$	$0.156$	Diverged	$33.4$ $\,\pm\,$	$39.3$	$36.0$ $\,\pm\,$	$65.9$	Diverged	$5.50$ $\,\pm\,$	$0.0997$
Torus (GMM)	$\mathbf{0.615}$ $\,\pm\,$	$0.0293$	Diverged	Diverged		Diverged		Diverged	$11.3$ $\,\pm\,$	$1.01$
Torus (Uniform)	$\mathbf{0.595}$ $\,\pm\,$	$0.0462$	Diverged	Diverged		Diverged		Diverged	$12.9$ $\,\pm\,$	$0.977$
Klein (GMM)	$\mathbf{0.626}$ $\,\pm\,$	$0.0280$	Diverged	Diverged		Diverged		Diverged	$10.5$ $\,\pm\,$	$0.729$
Klein (Uniform)	$\mathbf{0.644}$ $\,\pm\,$	$0.0624$	Diverged	Diverged		Diverged		Diverged	$11.6$ $\,\pm\,$	$0.890$

Appendix

Table 8: Baseline comparison measured by test-set unconditional path distribution loss

\mathcal{L}_{\mathrm{x}}^{\mathrm{traj}}

(

\downarrow

). Horizon:

3\,\mathrm{s}

Case	Ours		( Teng et al., 2026b )	( Li et al., 2020 )	( Li et al., 2020 ) $\star$		( Bartosh et al., 2025 )	( Kiyohara et al., 2025 )
b-ball (GMM)	$\mathbf{0.432}$ $\,\pm\,$	$0.00909$	Diverged	Diverged	$3.69$ $\,\pm\,$	$3.66$	Diverged	$0.730$ $\,\pm\,$	$0.00993$
b-ball (Uniform)	$\mathbf{0.338}$ $\,\pm\,$	$0.0119$	Diverged	Diverged	Diverged		Diverged	$0.534$ $\,\pm\,$	$0.00627$
Torus (GMM)	$\mathbf{0.687}$ $\,\pm\,$	$0.00109$	Diverged	Diverged	Diverged		Diverged	Diverged
Torus (Uniform)	$\mathbf{0.677}$ $\,\pm\,$	$0.00517$	Diverged	Diverged	Diverged		Diverged	$0.955$ $\,\pm\,$	$0.0207$
Klein (GMM)	$\mathbf{0.688}$ $\,\pm\,$	$0.00339$	Diverged	Diverged	Diverged		Diverged	$0.903$ $\,\pm\,$	$0.0164$
Klein (Uniform)	$\mathbf{0.676}$ $\,\pm\,$	$0.00832$	Diverged	Diverged	Diverged		Diverged	$0.922$ $\,\pm\,$	$0.0205$

Appendix

Table 9: Baseline comparison measured by test-set conditional path distribution loss

\mathcal{L}_{\mathrm{x}}^{\mathrm{traj,c}}

(

\downarrow

). Horizon:

5\,\mathrm{s}

Case	Ours		( Teng et al., 2026b )	( Li et al., 2020 )	( Li et al., 2020 ) $\star$		( Bartosh et al., 2025 )	( Kiyohara et al., 2025 )
b-ball (GMM)	$\mathbf{0.823}$ $\,\pm\,$	$0.144$	Diverged	Diverged	$154$ $\,\pm\,$	$187$	Diverged	$9.53$ $\,\pm\,$	$0.341$
b-ball (Uniform)	$\mathbf{0.709}$ $\,\pm\,$	$0.159$	Diverged	Diverged	Diverged		Diverged	$5.40$ $\,\pm\,$	$0.107$
Torus (GMM)	$\mathbf{0.795}$ $\,\pm\,$	$0.0315$	Diverged	Diverged	Diverged		Diverged	Diverged
Torus (Uniform)	$\mathbf{0.794}$ $\,\pm\,$	$0.0622$	Diverged	Diverged	Diverged		Diverged	$16.7$ $\,\pm\,$	$1.32$
Klein (GMM)	$\mathbf{0.833}$ $\,\pm\,$	$0.0324$	Diverged	Diverged	Diverged		Diverged	$13.4$ $\,\pm\,$	$1.04$
Klein (Uniform)	$\mathbf{0.826}$ $\,\pm\,$	$0.0696$	Diverged	Diverged	Diverged		Diverged	$14.7$ $\,\pm\,$	$1.27$

Appendix

Table 10: Baseline comparison measured by test-set unconditional path distribution loss

\mathcal{L}_{\mathrm{x}}^{\mathrm{traj}}

(

\downarrow

). Horizon:

5\,\mathrm{s}

	$d_{z}=d_{x}$		$d_{z}=2d_{x}$		$d_{z}=3d_{x}$		$d_{z}=4d_{x}$ (default)
b-ball (GMM)	$0.259$ $\,\pm\,$	$0.0172$	$0.275$ $\,\pm\,$	$0.0247$	$0.270$ $\,\pm\,$	$0.0227$	$\mathbf{0.257}$ $\,\pm\,$	$0.0158$
b-ball (Uniform)	$0.340$ $\,\pm\,$	$0.0712$	$0.338$ $\,\pm\,$	$0.0656$	$0.300$ $\,\pm\,$	$0.0212$	$\mathbf{0.289}$ $\,\pm\,$	$0.0317$
Torus (GMM)	$0.696$ $\,\pm\,$	$0.0440$	$0.614$ $\,\pm\,$	$0.0143$	$0.608$ $\,\pm\,$	$0.00535$	$\mathbf{0.603}$ $\,\pm\,$	$0.00347$
Torus (Uniform)	$0.685$ $\,\pm\,$	$0.0171$	$0.563$ $\,\pm\,$	$0.0121$	$\mathbf{0.553}$ $\,\pm\,$	$0.00406$	$0.560$ $\,\pm\,$	$0.00756$
Klein (GMM)	$0.689$ $\,\pm\,$	$0.0206$	$0.630$ $\,\pm\,$	$0.0189$	$\mathbf{0.611}$ $\,\pm\,$	$0.00336$	$0.612$ $\,\pm\,$	$0.0124$
Klein (Uniform)	$0.640$ $\,\pm\,$	$0.0354$	$0.562$ $\,\pm\,$	$0.00371$	$\mathbf{0.559}$ $\,\pm\,$	$0.00286$	$0.560$ $\,\pm\,$	$0.0178$

Appendix

Table 11: Latent-dimension ablation measured by test-set conditional path distribution loss

\mathcal{L}_{\mathrm{x}}^{\mathrm{traj,c}}

(

\downarrow

). Horizon:

1\,\mathrm{s}

	$d_{z}=d_{x}$		$d_{z}=2d_{x}$		$d_{z}=3d_{x}$		$d_{z}=4d_{x}$ (default)
b-ball (GMM)	$0.274$ $\,\pm\,$	$0.0452$	$0.311$ $\,\pm\,$	$0.107$	$0.336$ $\,\pm\,$	$0.147$	$\mathbf{0.214}$ $\,\pm\,$	$0.0558$
b-ball (Uniform)	$0.680$ $\,\pm\,$	$0.524$	$0.460$ $\,\pm\,$	$0.290$	$0.274$ $\,\pm\,$	$0.0762$	$\mathbf{0.249}$ $\,\pm\,$	$0.0819$
Torus (GMM)	$0.877$ $\,\pm\,$	$0.902$	$0.342$ $\,\pm\,$	$0.0256$	$0.350$ $\,\pm\,$	$0.0165$	$\mathbf{0.333}$ $\,\pm\,$	$0.0133$
Torus (Uniform)	$0.509$ $\,\pm\,$	$0.0629$	$0.302$ $\,\pm\,$	$0.0256$	$\mathbf{0.286}$ $\,\pm\,$	$0.0155$	$0.304$ $\,\pm\,$	$0.0319$
Klein (GMM)	$0.553$ $\,\pm\,$	$0.119$	$0.381$ $\,\pm\,$	$0.0473$	$\mathbf{0.339}$ $\,\pm\,$	$0.0272$	$0.341$ $\,\pm\,$	$0.0201$
Klein (Uniform)	$0.445$ $\,\pm\,$	$0.0601$	$0.296$ $\,\pm\,$	$0.0144$	$\mathbf{0.296}$ $\,\pm\,$	$0.0223$	$0.298$ $\,\pm\,$	$0.0317$

Appendix

Table 12: Latent-dimension ablation measured by test-set unconditional path distribution loss

\mathcal{L}_{\mathrm{x}}^{\mathrm{traj}}

(

\downarrow

). Horizon:

1\,\mathrm{s}

	DDIM				MLP
Case	Conditional		Unconditional		Conditional		Unconditional
b-ball (GMM)	$\mathbf{0.257}$ $\,\pm\,$	$0.0158$	$\mathbf{0.214}$ $\,\pm\,$	$0.0558$	$0.285$ $\,\pm\,$	$0.0162$	$0.344$ $\,\pm\,$	$0.0983$
b-ball (Uniform)	$\mathbf{0.289}$ $\,\pm\,$	$0.0317$	$\mathbf{0.249}$ $\,\pm\,$	$0.0819$	$0.332$ $\,\pm\,$	$0.0303$	$0.457$ $\,\pm\,$	$0.122$
Torus (GMM)	$\mathbf{0.603}$ $\,\pm\,$	$0.00347$	$\mathbf{0.333}$ $\,\pm\,$	$0.0133$	$0.610$ $\,\pm\,$	$0.00799$	$0.342$ $\,\pm\,$	$0.0214$
Torus (Uniform)	$0.560$ $\,\pm\,$	$0.00756$	$0.304$ $\,\pm\,$	$0.0319$	$\mathbf{0.560}$ $\,\pm\,$	$0.00377$	$\mathbf{0.299}$ $\,\pm\,$	$0.0136$
Klein (GMM)	$\mathbf{0.612}$ $\,\pm\,$	$0.0124$	$\mathbf{0.341}$ $\,\pm\,$	$0.0201$	$0.613$ $\,\pm\,$	$0.00592$	$0.345$ $\,\pm\,$	$0.00873$

Appendix

Table 13: Sampler ablation measured by test-set conditional and unconditional path distribution losses (

\downarrow

). Horizon:

1\,\mathrm{s}

Case	Default model		$\mathcal{L}_{\mathrm{x}}^{\mathrm{traj,c}}$		$\mathcal{L}_{z}$		$\mathcal{L}_{x}$		$\mathcal{L}_{\mathrm{x}}^{\mathrm{traj}}$		$\mathcal{L}_{\mathrm{x}}^{\mathrm{traj,c}},\ \mathcal{L}_{\mathrm{x}}^{\mathrm{traj}}$		$\mathcal{L}_{x},\ \mathcal{L}_{\mathrm{x}}^{\mathrm{traj}}$		$\mathcal{L}_{\mathrm{x}}^{\mathrm{traj,c}},\ \mathcal{L}_{x}$
b-ball (GMM)	$0.257$ $\,\pm\,$	$0.0158$	$0.255$ $\,\pm\,$	$0.00746$	$0.288$ $\,\pm\,$	$0.0348$	$0.280$ $\,\pm\,$	$0.0367$	$0.261$ $\,\pm\,$	$0.0208$	$0.270$ $\,\pm\,$	$0.0230$	$0.258$ $\,\pm\,$	$0.00631$	$\mathbf{0.254}$ $\,\pm\,$	$0.00987$
b-ball (Uniform)	$0.289$ $\,\pm\,$	$0.0317$	$0.292$ $\,\pm\,$	$0.0202$	$0.335$ $\,\pm\,$	$0.0489$	$0.297$ $\,\pm\,$	$0.0353$	$0.298$ $\,\pm\,$	$0.0346$	$0.317$ $\,\pm\,$	$0.0163$	$0.288$ $\,\pm\,$	$0.0150$	$\mathbf{0.283}$ $\,\pm\,$	$0.00573$
Torus (GMM)	$\mathbf{0.603}$ $\,\pm\,$	$0.00347$	$0.607$ $\,\pm\,$	$0.00351$	$0.649$ $\,\pm\,$	$0.0384$	$0.606$ $\,\pm\,$	$0.00463$	$0.617$ $\,\pm\,$	$0.00411$	$0.709$ $\,\pm\,$	$0.00138$	$0.616$ $\,\pm\,$	$0.00403$	$0.606$ $\,\pm\,$	$0.00167$
Torus (Uniform)	$0.560$ $\,\pm\,$	$0.00756$	$0.558$ $\,\pm\,$	$0.00334$	$0.647$ $\,\pm\,$	$0.0682$	$\mathbf{0.553}$ $\,\pm\,$	$0.00432$	$0.555$ $\,\pm\,$	$0.00296$	$0.709$ $\,\pm\,$	$0.00129$	$0.559$ $\,\pm\,$	$0.00389$	$0.554$ $\,\pm\,$	$0.00216$
Klein (GMM)	$0.612$ $\,\pm\,$	$0.0124$	$0.610$ $\,\pm\,$	$0.00808$	$0.709$ $\,\pm\,$	$0.0595$	$0.607$ $\,\pm\,$	$0.00778$	$0.614$ $\,\pm\,$	$0.0108$	$0.708$ $\,\pm\,$	$0.00172$	$0.616$ $\,\pm\,$	$0.00606$	$\mathbf{0.605}$ $\,\pm\,$	$0.00217$
Klein (Uniform)	$0.560$ $\,\pm\,$	$0.0178$	$0.559$ $\,\pm\,$	$0.00468$	$0.718$ $\,\pm\,$	$0.0175$	$0.562$ $\,\pm\,$	$0.0225$	$0.564$ $\,\pm\,$	$0.0176$	$0.704$ $\,\pm\,$	$0.000841$	$0.564$ $\,\pm\,$	$0.0108$	$\mathbf{0.554}$ $\,\pm\,$	$0.00259$

Appendix

Table 14: Loss-function ablation measured by test-set conditional path distribution loss

\mathcal{L}_{\mathrm{x}}^{\mathrm{traj,c}}

(

\downarrow

). Column labels other than Default model identify the removed objectives (their weights are set to zero). Values are the mean

\pm

sample standard deviation over five training seeds. The lowest mean among finite entries in each row is bold. Diverged denotes a configuration with nonfinite results or either test metric exceeding

10^{3}

in any seed. Horizon:

1\,\mathrm{s}

Case	Default model		$\mathcal{L}_{\mathrm{x}}^{\mathrm{traj,c}}$		$\mathcal{L}_{z}$		$\mathcal{L}_{x}$		$\mathcal{L}_{\mathrm{x}}^{\mathrm{traj}}$		$\mathcal{L}_{\mathrm{x}}^{\mathrm{traj,c}},\ \mathcal{L}_{\mathrm{x}}^{\mathrm{traj}}$		$\mathcal{L}_{x},\ \mathcal{L}_{\mathrm{x}}^{\mathrm{traj}}$		$\mathcal{L}_{\mathrm{x}}^{\mathrm{traj,c}},\ \mathcal{L}_{x}$
b-ball (GMM)	$0.214$ $\,\pm\,$	$0.0558$	$0.215$ $\,\pm\,$	$0.0395$	$0.411$ $\,\pm\,$	$0.273$	$0.379$ $\,\pm\,$	$0.247$	$\mathbf{0.194}$ $\,\pm\,$	$0.0268$	$0.308$ $\,\pm\,$	$0.157$	$0.216$ $\,\pm\,$	$0.0409$	$0.205$ $\,\pm\,$	$0.0457$
b-ball (Uniform)	$0.249$ $\,\pm\,$	$0.0819$	$0.279$ $\,\pm\,$	$0.110$	$0.408$ $\,\pm\,$	$0.169$	$0.260$ $\,\pm\,$	$0.113$	$0.268$ $\,\pm\,$	$0.107$	$0.335$ $\,\pm\,$	$0.0675$	$\mathbf{0.218}$ $\,\pm\,$	$0.0577$	$0.233$ $\,\pm\,$	$0.0397$
Torus (GMM)	$0.333$ $\,\pm\,$	$0.0133$	$0.323$ $\,\pm\,$	$0.00586$	$0.842$ $\,\pm\,$	$0.493$	$0.348$ $\,\pm\,$	$0.0204$	$0.386$ $\,\pm\,$	$0.0345$	$0.507$ $\,\pm\,$	$0.0134$	$0.381$ $\,\pm\,$	$0.0361$	$\mathbf{0.322}$ $\,\pm\,$	$0.0132$
Torus (Uniform)	$0.304$ $\,\pm\,$	$0.0319$	$0.285$ $\,\pm\,$	$0.00401$	$1.01$ $\,\pm\,$	$0.709$	$0.299$ $\,\pm\,$	$0.0205$	$0.302$ $\,\pm\,$	$0.0136$	$0.545$ $\,\pm\,$	$0.0112$	$0.322$ $\,\pm\,$	$0.0211$	$\mathbf{0.282}$ $\,\pm\,$	$0.00376$
Klein (GMM)	$0.341$ $\,\pm\,$	$0.0201$	$0.314$ $\,\pm\,$	$0.0113$	$1.67$ $\,\pm\,$	$1.02$	$0.339$ $\,\pm\,$	$0.0124$	$0.344$ $\,\pm\,$	$0.0334$	$0.515$ $\,\pm\,$	$0.00833$	$0.392$ $\,\pm\,$	$0.0368$	$\mathbf{0.313}$ $\,\pm\,$	$0.0125$
Klein (Uniform)	$0.298$ $\,\pm\,$	$0.0317$	$0.284$ $\,\pm\,$	$0.0140$	$1.60$ $\,\pm\,$	$0.489$	$0.300$ $\,\pm\,$	$0.0373$	$0.330$ $\,\pm\,$	$0.0412$	$0.544$ $\,\pm\,$	$0.00919$	$0.328$ $\,\pm\,$	$0.0268$	$\mathbf{0.282}$ $\,\pm\,$	$0.00749$

Appendix

Table 15: Loss-function ablation measured by test-set unconditional path distribution loss

\mathcal{L}_{\mathrm{x}}^{\mathrm{traj}}

(

\downarrow

). Column labels other than Default model identify the removed objectives (their weights are set to zero). Values are the mean

\pm

sample standard deviation over five training seeds. The lowest mean among finite entries in each row is bold. Diverged denotes a configuration with nonfinite results or either test metric exceeding

10^{3}

in any seed. Horizon:

1\,\mathrm{s}

Setting	$d_{z}=d_{x}$	$d_{z}=2d_{x}$	$d_{z}=3d_{x}$	MLP	$\mathcal{L}_{\mathrm{x}}^{\mathrm{traj,c}}$	$\mathcal{L}_{z}$	$\mathcal{L}_{x}$	$\mathcal{L}_{\mathrm{x}}^{\mathrm{traj}}$	$\mathcal{L}_{\mathrm{x}}^{\mathrm{traj,c}},\ \mathcal{L}_{\mathrm{x}}^{\mathrm{traj}}$	$\mathcal{L}_{x},\ \mathcal{L}_{\mathrm{x}}^{\mathrm{traj}}$	$\mathcal{L}_{\mathrm{x}}^{\mathrm{traj,c}},\ \mathcal{L}_{x}$
Mean change (%)	$+25.9\;(+417.1)$	$+1.5\;(+15.0)$	$+0.7\;(+14.4)$	$+2.9\;(+13.9)$	$+4.4\;(+2.6)$	$+63.8\;(+1611.9)$	$+0.5\;(+6.4)$	$-0.7\;(+2.7)$	$+12.5\;(+24.6)$	$-1.1\;(+5.1)$	$+3.9\;(-0.0)$
Worse cases (/11)	$11\;(11)$	$5\;(7)$	$5\;(6)$	$7\;(9)$	$8\;(5)$	$11\;(11)$	$6\;(7)$	$6\;(7)$	$11\;(11)$	$6\;(7)$	$7\;(4)$

Appendix

Table 16: Ablation summary across 11 benchmark settings on the conditional (unconditional) trajectory distribution loss. Mean change is the average of per-setting percentage changes relative to the default model; positive values indicate degradation. Worse cases count the settings with a higher mean test loss than the default model (out of 11), using the means over five training seeds. MLP replaces the diffusion sampler with a deterministic MLP encoder; loss labels identify the removed objectives. Per-case latent-dimension and loss-ablation results are reported in Tables 17 , 18 , 20 and 21 . Per-case conditional and unconditional results for the MLP encoder ablation are reported in Table 19 . Horizon:

3\,\mathrm{s}

	$d_{z}=d_{x}$		$d_{z}=2d_{x}$		$d_{z}=3d_{x}$		$d_{z}=4d_{x}$ (default)
b-ball (GMM)	$0.532$ $\,\pm\,$	$0.0294$	$0.505$ $\,\pm\,$	$0.0345$	$0.529$ $\,\pm\,$	$0.0523$	$\mathbf{0.487}$ $\,\pm\,$	$0.0153$
b-ball (Uniform)	$1.35$ $\,\pm\,$	$1.97$	$0.486$ $\,\pm\,$	$0.0746$	$0.438$ $\,\pm\,$	$0.0267$	$\mathbf{0.433}$ $\,\pm\,$	$0.0152$
Torus (GMM)	$0.722$ $\,\pm\,$	$0.0559$	$0.675$ $\,\pm\,$	$0.00497$	$0.673$ $\,\pm\,$	$0.00219$	$\mathbf{0.671}$ $\,\pm\,$	$0.00239$
Torus (Uniform)	$0.701$ $\,\pm\,$	$0.00576$	$0.652$ $\,\pm\,$	$0.00565$	$\mathbf{0.645}$ $\,\pm\,$	$0.00320$	$0.654$ $\,\pm\,$	$0.00691$
Klein (GMM)	$0.705$ $\,\pm\,$	$0.00810$	$0.683$ $\,\pm\,$	$0.00786$	$\mathbf{0.675}$ $\,\pm\,$	$0.00166$	$0.677$ $\,\pm\,$	$0.00530$
Klein (Uniform)	$0.686$ $\,\pm\,$	$0.0136$	$\mathbf{0.651}$ $\,\pm\,$	$0.00353$	$0.653$ $\,\pm\,$	$0.00316$	$0.653$ $\,\pm\,$	$0.0108$

Appendix

Table 17: Latent-dimension ablation measured by test-set conditional path distribution loss

\mathcal{L}_{\mathrm{x}}^{\mathrm{traj,c}}

(

\downarrow

). Horizon:

3\,\mathrm{s}

	$d_{z}=d_{x}$		$d_{z}=2d_{x}$		$d_{z}=3d_{x}$		$d_{z}=4d_{x}$ (default)
b-ball (GMM)	$1.39$ $\,\pm\,$	$0.375$	$1.03$ $\,\pm\,$	$0.377$	$1.37$ $\,\pm\,$	$0.709$	$\mathbf{0.732}$ $\,\pm\,$	$0.137$
b-ball (Uniform)	$28.4$ $\,\pm\,$	$60.3$	$1.18$ $\,\pm\,$	$0.838$	$\mathbf{0.644}$ $\,\pm\,$	$0.265$	$0.690$ $\,\pm\,$	$0.156$
Torus (GMM)	$2.17$ $\,\pm\,$	$3.18$	$0.623$ $\,\pm\,$	$0.0324$	$0.643$ $\,\pm\,$	$0.0275$	$\mathbf{0.615}$ $\,\pm\,$	$0.0293$
Torus (Uniform)	$0.785$ $\,\pm\,$	$0.0796$	$0.580$ $\,\pm\,$	$0.0555$	$\mathbf{0.560}$ $\,\pm\,$	$0.0218$	$0.595$ $\,\pm\,$	$0.0462$
Klein (GMM)	$0.897$ $\,\pm\,$	$0.167$	$0.689$ $\,\pm\,$	$0.0678$	$0.638$ $\,\pm\,$	$0.0407$	$\mathbf{0.626}$ $\,\pm\,$	$0.0280$
Klein (Uniform)	$0.786$ $\,\pm\,$	$0.0601$	$\mathbf{0.623}$ $\,\pm\,$	$0.0258$	$0.631$ $\,\pm\,$	$0.0214$	$0.644$ $\,\pm\,$	$0.0624$

Appendix

Table 18: Latent-dimension ablation measured by test-set unconditional path distribution loss

\mathcal{L}_{\mathrm{x}}^{\mathrm{traj}}

(

\downarrow

). Horizon:

3\,\mathrm{s}

	DDIM				MLP
Case	Conditional		Unconditional		Conditional		Unconditional
b-ball (GMM)	$\mathbf{0.487}$ $\,\pm\,$	$0.0153$	$\mathbf{0.732}$ $\,\pm\,$	$0.137$	$0.544$ $\,\pm\,$	$0.0305$	$1.43$ $\,\pm\,$	$0.570$
b-ball (Uniform)	$\mathbf{0.433}$ $\,\pm\,$	$0.0152$	$\mathbf{0.690}$ $\,\pm\,$	$0.156$	$0.474$ $\,\pm\,$	$0.0166$	$0.971$ $\,\pm\,$	$0.229$
Torus (GMM)	$\mathbf{0.671}$ $\,\pm\,$	$0.00239$	$\mathbf{0.615}$ $\,\pm\,$	$0.0293$	$0.674$ $\,\pm\,$	$0.00270$	$0.623$ $\,\pm\,$	$0.0338$
Torus (Uniform)	$0.654$ $\,\pm\,$	$0.00691$	$\mathbf{0.595}$ $\,\pm\,$	$0.0462$	$\mathbf{0.649}$ $\,\pm\,$	$0.00262$	$0.607$ $\,\pm\,$	$0.0386$
Klein (GMM)	$\mathbf{0.677}$ $\,\pm\,$	$0.00530$	$\mathbf{0.626}$ $\,\pm\,$	$0.0280$	$0.677$ $\,\pm\,$	$0.00337$	$0.644$ $\,\pm\,$	$0.0320$

Appendix

Table 19: Sampler ablation measured by test-set conditional and unconditional path distribution losses (

\downarrow

). Horizon:

3\,\mathrm{s}

Case	Default model		$\mathcal{L}_{\mathrm{x}}^{\mathrm{traj,c}}$		$\mathcal{L}_{z}$		$\mathcal{L}_{x}$		$\mathcal{L}_{\mathrm{x}}^{\mathrm{traj}}$		$\mathcal{L}_{\mathrm{x}}^{\mathrm{traj,c}},\ \mathcal{L}_{\mathrm{x}}^{\mathrm{traj}}$		$\mathcal{L}_{x},\ \mathcal{L}_{\mathrm{x}}^{\mathrm{traj}}$		$\mathcal{L}_{\mathrm{x}}^{\mathrm{traj,c}},\ \mathcal{L}_{x}$
b-ball (GMM)	$\mathbf{0.487}$ $\,\pm\,$	$0.0153$	$0.508$ $\,\pm\,$	$0.0103$	$0.588$ $\,\pm\,$	$0.0534$	$0.518$ $\,\pm\,$	$0.0343$	$0.498$ $\,\pm\,$	$0.0116$	$0.518$ $\,\pm\,$	$0.0274$	$0.497$ $\,\pm\,$	$0.0185$	$0.505$ $\,\pm\,$	$0.00621$
b-ball (Uniform)	$0.433$ $\,\pm\,$	$0.0152$	$0.446$ $\,\pm\,$	$0.0268$	$0.508$ $\,\pm\,$	$0.0556$	$0.436$ $\,\pm\,$	$0.0227$	$0.434$ $\,\pm\,$	$0.0140$	$0.456$ $\,\pm\,$	$0.0131$	$0.438$ $\,\pm\,$	$0.0137$	$\mathbf{0.432}$ $\,\pm\,$	$0.0140$
Torus (GMM)	$\mathbf{0.671}$ $\,\pm\,$	$0.00239$	$0.673$ $\,\pm\,$	$0.00159$	$0.731$ $\,\pm\,$	$0.0464$	$0.673$ $\,\pm\,$	$0.00224$	$0.678$ $\,\pm\,$	$0.00180$	$0.707$ $\,\pm\,$	$0.000957$	$0.677$ $\,\pm\,$	$0.00300$	$0.672$ $\,\pm\,$	$0.00146$
Torus (Uniform)	$0.654$ $\,\pm\,$	$0.00691$	$0.648$ $\,\pm\,$	$0.00229$	$0.730$ $\,\pm\,$	$0.0486$	$0.647$ $\,\pm\,$	$0.00285$	$0.647$ $\,\pm\,$	$0.00234$	$0.708$ $\,\pm\,$	$0.000933$	$0.652$ $\,\pm\,$	$0.00284$	$\mathbf{0.645}$ $\,\pm\,$	$0.00170$
Klein (GMM)	$0.677$ $\,\pm\,$	$0.00530$	$0.674$ $\,\pm\,$	$0.00478$	$0.813$ $\,\pm\,$	$0.0708$	$0.675$ $\,\pm\,$	$0.00368$	$0.678$ $\,\pm\,$	$0.00394$	$0.709$ $\,\pm\,$	$0.00120$	$0.679$ $\,\pm\,$	$0.000766$	$\mathbf{0.673}$ $\,\pm\,$	$0.00137$
Klein (Uniform)	$0.653$ $\,\pm\,$	$0.0108$	$0.651$ $\,\pm\,$	$0.00330$	$0.779$ $\,\pm\,$	$0.0504$	$0.652$ $\,\pm\,$	$0.0112$	$0.654$ $\,\pm\,$	$0.0104$	$0.706$ $\,\pm\,$	$0.000503$	$0.655$ $\,\pm\,$	$0.00693$	$\mathbf{0.647}$ $\,\pm\,$	$0.00151$

Appendix

Table 20: Loss-function ablation measured by test-set conditional path distribution loss

\mathcal{L}_{\mathrm{x}}^{\mathrm{traj,c}}

(

\downarrow

). Column labels other than Default model identify the removed objectives (their weights are set to zero). Values are the mean

\pm

sample standard deviation over five training seeds. The lowest mean among finite entries in each row is bold. Diverged denotes a configuration with nonfinite results or either test metric exceeding

10^{3}

in any seed. Horizon:

3\,\mathrm{s}

Case	Default model		$\mathcal{L}_{\mathrm{x}}^{\mathrm{traj,c}}$		$\mathcal{L}_{z}$		$\mathcal{L}_{x}$		$\mathcal{L}_{\mathrm{x}}^{\mathrm{traj}}$		$\mathcal{L}_{\mathrm{x}}^{\mathrm{traj,c}},\ \mathcal{L}_{\mathrm{x}}^{\mathrm{traj}}$		$\mathcal{L}_{x},\ \mathcal{L}_{\mathrm{x}}^{\mathrm{traj}}$		$\mathcal{L}_{\mathrm{x}}^{\mathrm{traj,c}},\ \mathcal{L}_{x}$
b-ball (GMM)	$\mathbf{0.732}$ $\,\pm\,$	$0.137$	$0.903$ $\,\pm\,$	$0.187$	$2.25$ $\,\pm\,$	$0.908$	$1.18$ $\,\pm\,$	$0.268$	$0.867$ $\,\pm\,$	$0.134$	$1.07$ $\,\pm\,$	$0.568$	$0.954$ $\,\pm\,$	$0.265$	$0.842$ $\,\pm\,$	$0.151$
b-ball (Uniform)	$0.690$ $\,\pm\,$	$0.156$	$0.646$ $\,\pm\,$	$0.216$	$1.62$ $\,\pm\,$	$1.19$	$0.696$ $\,\pm\,$	$0.176$	$0.717$ $\,\pm\,$	$0.0934$	$0.710$ $\,\pm\,$	$0.121$	$0.783$ $\,\pm\,$	$0.248$	$\mathbf{0.572}$ $\,\pm\,$	$0.150$
Torus (GMM)	$0.615$ $\,\pm\,$	$0.0293$	$\mathbf{0.603}$ $\,\pm\,$	$0.0168$	$2.81$ $\,\pm\,$	$1.70$	$0.651$ $\,\pm\,$	$0.0348$	$0.695$ $\,\pm\,$	$0.0517$	$0.741$ $\,\pm\,$	$0.0152$	$0.680$ $\,\pm\,$	$0.0673$	$0.604$ $\,\pm\,$	$0.0211$
Torus (Uniform)	$0.595$ $\,\pm\,$	$0.0462$	$0.551$ $\,\pm\,$	$0.00996$	$2.81$ $\,\pm\,$	$1.53$	$0.589$ $\,\pm\,$	$0.0381$	$0.586$ $\,\pm\,$	$0.0239$	$0.776$ $\,\pm\,$	$0.0203$	$0.633$ $\,\pm\,$	$0.0439$	$\mathbf{0.543}$ $\,\pm\,$	$0.00864$
Klein (GMM)	$0.626$ $\,\pm\,$	$0.0280$	$\mathbf{0.586}$ $\,\pm\,$	$0.0143$	$5.89$ $\,\pm\,$	$3.80$	$0.640$ $\,\pm\,$	$0.0394$	$0.654$ $\,\pm\,$	$0.0503$	$0.733$ $\,\pm\,$	$0.0108$	$0.703$ $\,\pm\,$	$0.0617$	$0.590$ $\,\pm\,$	$0.0167$
Klein (Uniform)	$0.644$ $\,\pm\,$	$0.0624$	$0.618$ $\,\pm\,$	$0.0222$	$3.90$ $\,\pm\,$	$1.99$	$0.643$ $\,\pm\,$	$0.0678$	$0.684$ $\,\pm\,$	$0.0529$	$0.841$ $\,\pm\,$	$0.0198$	$0.689$ $\,\pm\,$	$0.0346$	$\mathbf{0.608}$ $\,\pm\,$	$0.0166$

Appendix

Table 21: Loss-function ablation measured by test-set unconditional path distribution loss

\mathcal{L}_{\mathrm{x}}^{\mathrm{traj}}

(

\downarrow

). Column labels other than Default model identify the removed objectives (their weights are set to zero). Values are the mean

\pm

sample standard deviation over five training seeds. The lowest mean among finite entries in each row is bold. Diverged denotes a configuration with nonfinite results or either test metric exceeding

10^{3}

in any seed. Horizon:

3\,\mathrm{s}

Setting	$d_{z}=d_{x}$	$d_{z}=2d_{x}$	$d_{z}=3d_{x}$	MLP	$\mathcal{L}_{\mathrm{x}}^{\mathrm{traj,c}}$	$\mathcal{L}_{z}$	$\mathcal{L}_{x}$	$\mathcal{L}_{\mathrm{x}}^{\mathrm{traj}}$	$\mathcal{L}_{\mathrm{x}}^{\mathrm{traj,c}},\ \mathcal{L}_{\mathrm{x}}^{\mathrm{traj}}$	$\mathcal{L}_{x},\ \mathcal{L}_{\mathrm{x}}^{\mathrm{traj}}$	$\mathcal{L}_{\mathrm{x}}^{\mathrm{traj,c}},\ \mathcal{L}_{x}$
Mean change (%)	$\text{Diverged}\;(\text{Diverged})$	$\text{Diverged}\;(\text{Diverged})$	$+20.9\;(+396.2)$	$+3.1\;(+14.2)$	$+3.7\;(+1.6)$	$\text{Diverged}\;(\text{Diverged})$	$+0.8\;(+11.6)$	$-0.4\;(+4.0)$	$+9.9\;(+17.4)$	$-1.0\;(+4.8)$	$+3.2\;(-1.1)$
Worse cases (/11)	$\text{Diverged}\;(\text{Diverged})$	$\text{Diverged}\;(\text{Diverged})$	$7\;(5)$	$7\;(10)$	$8\;(4)$	$\text{Diverged}\;(\text{Diverged})$	$7\;(7)$	$6\;(6)$	$11\;(11)$	$6\;(8)$	$8\;(4)$

Appendix

Table 22: Ablation summary across 11 benchmark settings on the conditional (unconditional) trajectory distribution loss. Mean change is the average of per-setting percentage changes relative to the default model; positive values indicate degradation. Worse cases count the settings with a higher mean test loss than the default model (out of 11), using the means over five training seeds. MLP replaces the diffusion sampler with a deterministic MLP encoder; loss labels identify the removed objectives. Per-case latent-dimension and loss-ablation results are reported in Tables 23 , 24 , 26 and 27 . Per-case conditional and unconditional results for the MLP encoder ablation are reported in Table 25 . Horizon:

5\,\mathrm{s}

	$d_{z}=d_{x}$		$d_{z}=2d_{x}$		$d_{z}=3d_{x}$		$d_{z}=4d_{x}$ (default)
b-ball (GMM)	$0.469$ $\,\pm\,$	$0.0170$	$0.444$ $\,\pm\,$	$0.0239$	$0.465$ $\,\pm\,$	$0.0384$	$\mathbf{0.432}$ $\,\pm\,$	$0.00909$
b-ball (Uniform)	Diverged		$0.381$ $\,\pm\,$	$0.0576$	$0.343$ $\,\pm\,$	$0.0192$	$\mathbf{0.338}$ $\,\pm\,$	$0.0119$
Torus (GMM)	$0.771$ $\,\pm\,$	$0.153$	$0.689$ $\,\pm\,$	$0.00293$	$0.689$ $\,\pm\,$	$0.00180$	$\mathbf{0.687}$ $\,\pm\,$	$0.00109$
Torus (Uniform)	$0.705$ $\,\pm\,$	$0.00415$	$0.674$ $\,\pm\,$	$0.00480$	$\mathbf{0.671}$ $\,\pm\,$	$0.00285$	$0.677$ $\,\pm\,$	$0.00517$
Klein (GMM)	$0.707$ $\,\pm\,$	$0.00579$	$0.693$ $\,\pm\,$	$0.00565$	$0.688$ $\,\pm\,$	$0.000717$	$\mathbf{0.688}$ $\,\pm\,$	$0.00339$
Klein (Uniform)	$0.694$ $\,\pm\,$	$0.00823$	$\mathbf{0.675}$ $\,\pm\,$	$0.00388$	$0.677$ $\,\pm\,$	$0.00225$	$0.676$ $\,\pm\,$	$0.00832$

Appendix

Table 23: Latent-dimension ablation measured by test-set conditional path distribution loss

\mathcal{L}_{\mathrm{x}}^{\mathrm{traj,c}}

(

\downarrow

). Horizon:

5\,\mathrm{s}

	$d_{z}=d_{x}$		$d_{z}=2d_{x}$		$d_{z}=3d_{x}$		$d_{z}=4d_{x}$ (default)
b-ball (GMM)	$1.63$ $\,\pm\,$	$0.465$	$1.16$ $\,\pm\,$	$0.401$	$1.53$ $\,\pm\,$	$0.797$	$\mathbf{0.823}$ $\,\pm\,$	$0.144$
b-ball (Uniform)	Diverged		$1.17$ $\,\pm\,$	$0.816$	$\mathbf{0.654}$ $\,\pm\,$	$0.254$	$0.709$ $\,\pm\,$	$0.159$
Torus (GMM)	$6.81$ $\,\pm\,$	$13.2$	$0.801$ $\,\pm\,$	$0.0320$	$0.825$ $\,\pm\,$	$0.0328$	$\mathbf{0.795}$ $\,\pm\,$	$0.0315$
Torus (Uniform)	$0.962$ $\,\pm\,$	$0.0922$	$0.769$ $\,\pm\,$	$0.0617$	$\mathbf{0.761}$ $\,\pm\,$	$0.0301$	$0.794$ $\,\pm\,$	$0.0622$
Klein (GMM)	$1.13$ $\,\pm\,$	$0.199$	$0.890$ $\,\pm\,$	$0.0836$	$\mathbf{0.830}$ $\,\pm\,$	$0.0552$	$0.833$ $\,\pm\,$	$0.0324$
Klein (Uniform)	$0.959$ $\,\pm\,$	$0.0524$	$\mathbf{0.792}$ $\,\pm\,$	$0.0185$	$0.807$ $\,\pm\,$	$0.0245$	$0.826$ $\,\pm\,$	$0.0696$

Appendix

Table 24: Latent-dimension ablation measured by test-set unconditional path distribution loss

\mathcal{L}_{\mathrm{x}}^{\mathrm{traj}}

(

\downarrow

). Horizon:

5\,\mathrm{s}

	DDIM				MLP
Case	Conditional		Unconditional		Conditional		Unconditional
b-ball (GMM)	$\mathbf{0.432}$ $\,\pm\,$	$0.00909$	$\mathbf{0.823}$ $\,\pm\,$	$0.144$	$0.469$ $\,\pm\,$	$0.0210$	$1.48$ $\,\pm\,$	$0.567$
b-ball (Uniform)	$\mathbf{0.338}$ $\,\pm\,$	$0.0119$	$\mathbf{0.709}$ $\,\pm\,$	$0.159$	$0.370$ $\,\pm\,$	$0.0139$	$0.978$ $\,\pm\,$	$0.222$
Torus (GMM)	$\mathbf{0.687}$ $\,\pm\,$	$0.00109$	$\mathbf{0.795}$ $\,\pm\,$	$0.0315$	$0.689$ $\,\pm\,$	$0.00147$	$0.815$ $\,\pm\,$	$0.0344$
Torus (Uniform)	$0.677$ $\,\pm\,$	$0.00517$	$\mathbf{0.794}$ $\,\pm\,$	$0.0622$	$\mathbf{0.673}$ $\,\pm\,$	$0.00316$	$0.815$ $\,\pm\,$	$0.0550$
Klein (GMM)	$\mathbf{0.688}$ $\,\pm\,$	$0.00339$	$\mathbf{0.833}$ $\,\pm\,$	$0.0324$	$0.690$ $\,\pm\,$	$0.00186$	$0.840$ $\,\pm\,$	$0.0343$

Appendix

Table 25: Sampler ablation measured by test-set conditional and unconditional path distribution losses (

\downarrow

). Horizon:

5\,\mathrm{s}

Case	Default model		$\mathcal{L}_{\mathrm{x}}^{\mathrm{traj,c}}$		$\mathcal{L}_{z}$		$\mathcal{L}_{x}$		$\mathcal{L}_{\mathrm{x}}^{\mathrm{traj}}$		$\mathcal{L}_{\mathrm{x}}^{\mathrm{traj,c}},\ \mathcal{L}_{\mathrm{x}}^{\mathrm{traj}}$		$\mathcal{L}_{x},\ \mathcal{L}_{\mathrm{x}}^{\mathrm{traj}}$		$\mathcal{L}_{\mathrm{x}}^{\mathrm{traj,c}},\ \mathcal{L}_{x}$
b-ball (GMM)	$\mathbf{0.432}$ $\,\pm\,$	$0.00909$	$0.446$ $\,\pm\,$	$0.00769$	$0.512$ $\,\pm\,$	$0.0405$	$0.456$ $\,\pm\,$	$0.0193$	$0.442$ $\,\pm\,$	$0.0118$	$0.454$ $\,\pm\,$	$0.0190$	$0.444$ $\,\pm\,$	$0.0144$	$0.445$ $\,\pm\,$	$0.00619$
b-ball (Uniform)	$\mathbf{0.338}$ $\,\pm\,$	$0.0119$	$0.348$ $\,\pm\,$	$0.0204$	$0.403$ $\,\pm\,$	$0.0524$	$0.342$ $\,\pm\,$	$0.0167$	$0.342$ $\,\pm\,$	$0.0113$	$0.356$ $\,\pm\,$	$0.00949$	$0.342$ $\,\pm\,$	$0.00911$	$0.339$ $\,\pm\,$	$0.00832$
Torus (GMM)	$\mathbf{0.687}$ $\,\pm\,$	$0.00109$	$0.688$ $\,\pm\,$	$0.00111$	$0.755$ $\,\pm\,$	$0.0550$	$0.689$ $\,\pm\,$	$0.00123$	$0.691$ $\,\pm\,$	$0.000986$	$0.709$ $\,\pm\,$	$0.000482$	$0.690$ $\,\pm\,$	$0.00178$	$0.688$ $\,\pm\,$	$0.000696$
Torus (Uniform)	$0.677$ $\,\pm\,$	$0.00517$	$0.673$ $\,\pm\,$	$0.00130$	$0.748$ $\,\pm\,$	$0.0434$	$\mathbf{0.670}$ $\,\pm\,$	$0.00201$	$0.671$ $\,\pm\,$	$0.00221$	$0.708$ $\,\pm\,$	$0.000951$	$0.675$ $\,\pm\,$	$0.00248$	$0.670$ $\,\pm\,$	$0.00107$
Klein (GMM)	$0.688$ $\,\pm\,$	$0.00339$	$0.687$ $\,\pm\,$	$0.00281$	$0.897$ $\,\pm\,$	$0.208$	$0.689$ $\,\pm\,$	$0.00249$	$0.690$ $\,\pm\,$	$0.00326$	$0.707$ $\,\pm\,$	$0.000818$	$0.690$ $\,\pm\,$	$0.000768$	$\mathbf{0.685}$ $\,\pm\,$	$0.000721$
Klein (Uniform)	$0.676$ $\,\pm\,$	$0.00832$	$0.674$ $\,\pm\,$	$0.00120$	$0.795$ $\,\pm\,$	$0.0644$	$0.675$ $\,\pm\,$	$0.00772$	$0.677$ $\,\pm\,$	$0.00610$	$0.707$ $\,\pm\,$	$0.00122$	$0.678$ $\,\pm\,$	$0.00651$	$\mathbf{0.671}$ $\,\pm\,$	$0.000704$

Appendix

Table 26: Loss-function ablation measured by test-set conditional path distribution loss

\mathcal{L}_{\mathrm{x}}^{\mathrm{traj,c}}

(

\downarrow

). Column labels other than Default model identify the removed objectives (their weights are set to zero). Values are the mean

\pm

sample standard deviation over five training seeds. The lowest mean among finite entries in each row is bold. Diverged denotes a configuration with nonfinite results or either test metric exceeding

10^{3}

in any seed. Horizon:

5\,\mathrm{s}

Case	Default model		$\mathcal{L}_{\mathrm{x}}^{\mathrm{traj,c}}$		$\mathcal{L}_{z}$		$\mathcal{L}_{x}$		$\mathcal{L}_{\mathrm{x}}^{\mathrm{traj}}$		$\mathcal{L}_{\mathrm{x}}^{\mathrm{traj,c}},\ \mathcal{L}_{\mathrm{x}}^{\mathrm{traj}}$		$\mathcal{L}_{x},\ \mathcal{L}_{\mathrm{x}}^{\mathrm{traj}}$		$\mathcal{L}_{\mathrm{x}}^{\mathrm{traj,c}},\ \mathcal{L}_{x}$
b-ball (GMM)	$\mathbf{0.823}$ $\,\pm\,$	$0.144$	$0.990$ $\,\pm\,$	$0.219$	$2.59$ $\,\pm\,$	$1.06$	$1.29$ $\,\pm\,$	$0.188$	$1.09$ $\,\pm\,$	$0.295$	$1.19$ $\,\pm\,$	$0.538$	$1.18$ $\,\pm\,$	$0.347$	$0.931$ $\,\pm\,$	$0.191$
b-ball (Uniform)	$0.709$ $\,\pm\,$	$0.159$	$0.653$ $\,\pm\,$	$0.207$	$2.05$ $\,\pm\,$	$1.90$	$0.699$ $\,\pm\,$	$0.177$	$0.738$ $\,\pm\,$	$0.107$	$0.711$ $\,\pm\,$	$0.117$	$0.793$ $\,\pm\,$	$0.232$	$\mathbf{0.593}$ $\,\pm\,$	$0.155$
Torus (GMM)	$0.795$ $\,\pm\,$	$0.0315$	$\mathbf{0.770}$ $\,\pm\,$	$0.0216$	$4.45$ $\,\pm\,$	$2.87$	$0.835$ $\,\pm\,$	$0.0448$	$0.883$ $\,\pm\,$	$0.0634$	$0.899$ $\,\pm\,$	$0.0177$	$0.869$ $\,\pm\,$	$0.0769$	$0.782$ $\,\pm\,$	$0.0221$
Torus (Uniform)	$0.794$ $\,\pm\,$	$0.0622$	$0.732$ $\,\pm\,$	$0.00632$	$4.08$ $\,\pm\,$	$2.07$	$0.789$ $\,\pm\,$	$0.0612$	$0.778$ $\,\pm\,$	$0.0321$	$0.924$ $\,\pm\,$	$0.0233$	$0.826$ $\,\pm\,$	$0.0568$	$\mathbf{0.726}$ $\,\pm\,$	$0.00525$
Klein (GMM)	$0.833$ $\,\pm\,$	$0.0324$	$\mathbf{0.771}$ $\,\pm\,$	$0.0142$	$12.7$ $\,\pm\,$	$13.8$	$0.828$ $\,\pm\,$	$0.0536$	$0.844$ $\,\pm\,$	$0.0674$	$0.900$ $\,\pm\,$	$0.0208$	$0.898$ $\,\pm\,$	$0.0789$	$0.780$ $\,\pm\,$	$0.0140$
Klein (Uniform)	$0.826$ $\,\pm\,$	$0.0696$	$0.789$ $\,\pm\,$	$0.0250$	$5.78$ $\,\pm\,$	$3.49$	$0.829$ $\,\pm\,$	$0.0866$	$0.877$ $\,\pm\,$	$0.0660$	$0.991$ $\,\pm\,$	$0.0186$	$0.872$ $\,\pm\,$	$0.0374$	$\mathbf{0.786}$ $\,\pm\,$	$0.0125$

Appendix

Table 27: Loss-function ablation measured by test-set unconditional path distribution loss

\mathcal{L}_{\mathrm{x}}^{\mathrm{traj}}

(

\downarrow

). Column labels other than Default model identify the removed objectives (their weights are set to zero). Values are the mean

\pm

sample standard deviation over five training seeds. The lowest mean among finite entries in each row is bold. Diverged denotes a configuration with nonfinite results or either test metric exceeding

10^{3}

in any seed. Horizon:

5\,\mathrm{s}

Figure 7: b-ball (GMM): one-dimensional bouncing ball with GMM-distributed restitution coefficients. The two columns show position and velocity over time.

Figure 8: b-ball (Uniform): one-dimensional bouncing ball with uniformly distributed restitution coefficients. The two columns show position and velocity over time.

Figure 9: Torus (GMM): coordinate trajectories in the fundamental domain with GMM-distributed reset perturbations.

Figure 10: Torus (Uniform): coordinate trajectories in the fundamental domain with uniformly distributed reset perturbations.

Figure 11: Klein (GMM): coordinate trajectories in the fundamental domain with GMM-distributed reset perturbations and Klein-bottle boundary gluing.

Figure 12: Klein (Uniform): coordinate trajectories in the fundamental domain with uniformly distributed reset perturbations and Klein-bottle boundary gluing.

Figure 13: Klein–Torus: coordinate trajectories with random selection between canonical Klein-bottle and torus resets at the top boundary, each with probability

1/2

Figure 14: 1-ball (GMM): a single planar bouncing ball with GMM-distributed wall and floor restitution coefficients. The left

2\times 2

panels show

(x,y)

above

(v_{x},v_{y})

; the right panel shows the spatial path with radius-adjusted ball-center contact boundaries.

Figure 15: 1-ball (Uniform): a single planar bouncing ball with uniformly distributed wall and floor restitution coefficients. The left panels show the four state components over time; the right panel shows the spatial path and ball-center contact boundaries.

Figure 16: 2-balls (GMM): two interacting planar bouncing balls with GMM-distributed wall and floor restitution coefficients. The left panels show the first ball’s four state components; the two right panels show the spatial paths of both balls and their ball-center contact boundaries.

Figure 17: 2-balls (Uniform): two interacting planar bouncing balls with uniformly distributed wall and floor restitution coefficients. The left panels show the first ball’s four state components; the two right panels show the spatial paths of both balls and their ball-center contact boundaries.

Explore similar work

Jun 9, 2026cs.LG

Embedding Hybrid Systems into Continuous Latent Vector Fields

This work proves that an

n

-dimensional hybrid system can be embedded into an

m

-dimensional Euclidean space equipped with a continuous vector field on its embedded image whenever

m>2n

. This result suggests that an intrinsically discontinuous hybrid system generically admits a continuous extrinsic representation that is well-posed for differentiable optimization. Building on this existence theorem, we show that a latent Neural ODE with consistency loss in both the latent and state space can accurately recover the flow of hybrid systems. Extensive experiments suggest the proposed method outperforms the existing method in learning hybrid systems with varying geometries from only time series data.

Sangli Teng, Hang Liu, Koushil Sreenath

University of California, Berkeley, CA, USA · University of Michigan, Ann Arbor, MI, USA

Neural Ordinary Differential Equations Latent Dynamics Latent Space Differentiable Optimization

Ask PDF

Jun 15, 2026stat.ML

Closing the Approximation Gap in Simulation-free Latent SDEs

Recovering dynamical systems from noisy observations is a recurring challenge across scientific domains, including neuroscience and physics. Latent stochastic differential equations (SDEs) address this by modeling the system as an unobserved state that evolves according to a learnable SDE and generates the observations. Variational inference (VI) provides a tractable objective for fitting latent SDEs. Traditional VI algorithms evaluate this objective by numerical simulation over a time discretization, trading fidelity for computational cost. A recent class of algorithms, simulation-free VI, sidesteps this tradeoff by parameterizing the posterior through its instantaneous marginals rather than its drift. In this work, we show that the efficiency of existing simulation-free VI algorithms comes at a price: their parameterizations restrict the approximate posterior to a subset of the SDEs available to simulation-based methods, degrading posterior inference and parameter learning. We propose Helmholtz-SDE, a simulation-free VI algorithm that closes this gap by optimizing over path laws compatible with a prescribed collection of marginals. Helmholtz-SDE recovers dynamics more faithfully than prior simulation-free methods, with the largest gains under high posterior uncertainty. It further matches the performance of simulation-based VI at a fraction of the runtime. Code is available at https://github.com/lindermanlab/helmholtz-sde .

Henry D. Smith, Brian L. Trippe, Scott W. Linderman

Stanford University

Stochastic Differential Equations Variational Inference Latent Dynamics Closure Approximation Algorithms

Ask PDF

Jul 21, 2026cs.LG

Neural Kolmogorov Equations: Parallelizable Learning of Stochastic Dynamics under General Noise

Neural stochastic differential equations (SDEs) have emerged as powerful tools for learning noisy or stochastic dynamics directly from data; however, existing approaches largely assume uncoupled and continuous noise, limiting their applicability to realistic stochastic drivers, and often scale poorly in time, requiring expensive autoregressive training. To address these limitations, we propose Neural Kolmogorov Equations (NKEs), a deterministic, infinite-dimensional reformulation of Neural SDEs based on the Kolmogorov Forward equation, transforming the learning problem from modelling individual stochastic trajectories to modelling the evolution of probability densities. NKEs learn general Lévy-type stochastic forcing directly through the operator structure of the KFE, and enable parallel-in-time training via a Lagrangian Galerkin projection and operator splitting. We evaluate NKEs on several stochastic benchmarks, including systems with coupled noise and jump processes, and verify that NKEs provide flexible models that accurately recover deterministic and stochastic dynamics with competitive predictive accuracy and improved training efficiency. Code and pretrained models will be released.

Arthur Bizzi, Olga Fink

EPFL, Imperial College · EPFL

Stochastic Differential Equations Stochastic Dynamics Kolmogorov-Arnold Networks

Ask PDF