Generator-Flow Equivariance Training

Insert the method into a latent predictor or world model. Encode x_t as z_t=E_theta(x_t), predict z_hat_(t+1)=P_theta(z_t), and obtain generators G_a from the learned Lie-algebra module. For each minibatch, select generator indices and amplitudes epsilon in {epsilon_0,2 epsilon_0}. Compute Phi(z)=z+epsilon G_a(z), evaluate P_theta(Phi(z)), evaluate Phi(P_theta(z)), and add their squared difference to the prediction loss. A more efficient infinitesimal implementation computes the commutator with two Jacobian-vector products and uses L_inf. Use several small epsilon values and stop gradients through G during early training so the system cannot satisfy equivariance by making generators zero. The learned-generator construction is taken from the paper; finite-flow augmentation and predictor commutation are the neural-network adaptation. First test on a latent predictor trained on rotated 2D damped-oscillator or isotropic Ornstein-Uhlenbeck trajectories, withholding several rotation angles during training. Compare ordinary prediction, random augmentation, and learned-flow equivariance. Plot equivariance error against epsilon on log-log axes. The quantitative prediction is that the squared error scales as epsilon^4 after first-order equivariance is learned, versus epsilon^2 for a non-equivariant baseline; rotated-test error should decrease when normalized commutator residual crosses below 10^-2.