Generative World Models to compute protein folding pathways

We developed a generative World Model in the context of biomolecular simulations for proteins. An agent was trained by evolution strategy in a latent spatio-temporal representation of the structural transitions, to learn a policy to drive protein folding simulations via dihedral rotations (Φ, Ψ). Latent configurations were decoded back into atomistic structures using atomistic force fields, allowing for the calculation of structural properties and energetic terms. We show how the model can be regularized by adding Ramachandran-based rewards during the training of the controller. Results have been validated against equilibrium MD data. Markov State Modeling could be applied to reconstruct the dynamics of the system and reconstruct the unbiased folding landscape. The method proposed here facilitates the study of the conformational changes and folding pathways of proteins by generative AI, and can have application in structure-based drug discovery.