waxmorph.torch.train.TrainConfig

Contents

waxmorph.torch.train.TrainConfig#

class waxmorph.torch.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)[source]#

Bases: object

Rollout and optimizer settings for a fixed agent count.

n_epochs counts optimizer updates; histories evaluate each resulting model, requiring one extra rollout. Each rollout has t_rollout learned steps. mech_steps and diff_steps are prescribed substep counts; zero selects the identity path for that correction. dt_gns scales all learned deltas. dt_mech and dt_diff are explicit step sizes whose stable range depends on the active forces and graph; nonfinite states are rejected. Diffusion applies \(c\leftarrow\max(c-D_{emu}L_Gc\,\Delta t_{diff},0)\).

The optimized loss is \(L=L_{shape}+\lambda_{reg}L_{reg}\), where

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

The regularizer measures learned displacement before the mechanics correction. grad_clip_norm caps the global L2 gradient norm; None keeps the computed gradients. log_every controls progress output. JAX adds max_edges_factor for its static edge capacity.

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:
Contents