waxmorph.jax.train.TrainConfig

class waxmorph.jax.train.TrainConfig(n_epochs=2000, t_rollout=100, mech_steps=5, diff_steps=5, dt_mech=0.01, dt_diff=0.01, dt_gns=0.01, D_emu=0.1, lambda_reg=0.001, grad_clip_norm=1.0, log_every=10, max_edges_factor=10)

Bases: object

JAX rollout and optimizer settings for a fixed agent count.

n_epochs counts optimizer updates. Histories evaluate every post-update model, so training performs one additional rollout. Zero mech_steps or diff_steps selects the identity path for that prescribed update. dt_mech and dt_diff are explicit-Euler steps; diffusion stability depends on D_emu * dt_diff and graph degree. dt_gns scales the raw dX, dP, and dc heads before state updates.

The total loss is

\[L=L_{shape}+\lambda_{reg}\sum_t \lVert \Delta t_{GNS}\,GNS_{X,t}\rVert_F^2,\]

where the displacement term is measured before mechanics. grad_clip_norm=None keeps the computed gradients. JAX adds max_edges_factor: both directed-edge and neighbor-pair buffers are bounded by N * max_edges_factor; overflow raises while preserving the configured static shape.

Examples

>>> cfg = TrainConfig(n_epochs=3, t_rollout=2)
>>> print(cfg.n_epochs, cfg.t_rollout, cfg.grad_clip_norm)
3 2 1.0

Parameters:

• n_epochs (int)

• t_rollout (int)

• mech_steps (int)

• diff_steps (int)

• dt_mech (float)

• dt_diff (float)

• dt_gns (float)

• D_emu (float)

• lambda_reg (float)

• grad_clip_norm (float | None)

• log_every (int)

• max_edges_factor (int)