Constrained Graph Mechanics Networks

Learning to reason about relations and dynamics over multiple interacting objects is a prevailing yet challenging topic in machine learning. The challenges mainly stem from that interacting systems are of combinatorial complexity, the model we learn should obey Euclidean equivariance, and certain geometrical constraints commonly exist. Current methods, particularly the ones based on equivariant Graph Neural Networks (GNNs), have targeted on the first two challenges but remain immature for constrained systems. In this paper, we propose Graph Mechanics Network (GMN) which is combinatorially efficient, equivariant and constraint-aware. The core of GMN is that it represents, by generalized coordinates, the forward kinematics information (positions and velocities) of a structural object. In this manner, the geometrical constraints are implicitly and naturally encoded in the forward kinematics. Moreover, to allow equivariant message passing in GMN, we have developed a general form of orthogonality-equivariant functions, given that the dynamics of constrained systems are more complicated than the unconstrained counterparts. Theoretically, the proposed equivariant formulation is proved to be universally expressive under certain conditions. Simulated experiments on particles, sticks, hinges, and the combinations of them support the advantages of GMN compared to the state-of-the-art GNNs in terms of prediction accuracy, constraint satisfaction and data efficiency.