waxmorph.jax.train.train

Contents

waxmorph.jax.train.train#

waxmorph.jax.train.train(model, optimizer, opt_state, loss_fn, *, source_pos, polarities, c, radii, targets=None, config=None, save_path=None, device='cuda')[source]#

Train a GNS model for shape assembly with a fixed agent count (JAX backend).

Each epoch performs one update. Histories, the selected model, and its trajectory all describe post-update states; evaluating n_epochs updates requires n_epochs + 1 collect-and-replay rollouts. Collection freezes data-dependent topology outside autodiff; replay differentiates features and state updates through fixed, padded buffers.

Inputs are validated before conversion to contiguous float32. State arrays share a nonzero particle axis. Source polarities are used as supplied for the first graph, whose angle feature assumes unit vectors; later learned updates apply normalization. Target frame 0 means the state after the first rollout update, and valid unique frames lie in [0, t_rollout). TrainResult contains the selected model and log; callers manage Optax state separately.

A new save_path receives the best model’s two-file checkpoint and log. Use checkpoints from trusted sources; config, PyTree, leaf shapes, and dtypes must match. Saved weights replace the supplied model, Optax state is reinitialized, and exactly one refinement update preserves existing files.

Warp mechanics or diffusion uses custom VJPs on a CUDA-backed JAX device. CPU training uses zero for both prescribed step counts. See TrainResult for the log schema.

Raises:

  • TypeError – If config types, array dtypes, or target frames are invalid.

  • RuntimeError – If requested Warp physics fails CUDA-backed JAX/Warp validation.

  • ValueError – If config bounds, state arrays, radii, targets, or checkpoint structure are invalid.

Parameters:

  • model (JaxGNS)

  • optimizer (optax.GradientTransformation)

  • opt_state (optax.OptState)

  • loss_fn (Callable[[jnp.ndarray, jnp.ndarray], jnp.ndarray])

  • source_pos (np.ndarray)

  • polarities (np.ndarray)

  • c (np.ndarray)

  • radii (np.ndarray)

  • targets (list[tuple[int, np.ndarray]] | None)

  • config (TrainConfig | None)

  • save_path (str | Path | None)

  • device (str)

Return type:

TrainResult

Contents