# Rigid Body Dynamics Simulation Based on GNNs with Constraints



In recent years, the utilization of Graph Neural Network (GNN)-based methods for simulating complex physical systems has significantly advanced the field of numerical science and engineering. Despite their success, current GNN-based methods for rigid body dynamic simulation are constrained to relatively simple scenarios, hindering their practical use in industrial settings where complex mechanical structures and interconnected components prevail. These methods face challenges in handling intricate force relationships within rigid bodies, primarily due to the difficulty in obtaining force-related data for objects in industrial environments. To address this, we propose a novel constraint-guided method that incorporates force computations into GNN-based simulations. The model incorporates computations related to both contact and non-contact forces into the prediction process. Additionally, it imposes physical constraints on the prediction process based on Kane's equations. We have rigorously demonstrated the model's rationality and effectiveness with thorough theoretical demonstration and empirical analysis.





## Dependencies

```
python==3.8
scikit-learn==0.24.2
networkx==2.5.1
torch==1.8.0
torch-geometric==2.0.1
```


## Training and Evaluation

Firstly, make sure the dataset is ready and placed as instructed by the dataset. At this time, all placeholders like *YOUR_DATA_DIR* in the code here in this repository should be properly replaced as your data directory for the preprocessed data. The data is provided within the dataset folder.

*1. Training.* 

In the Physion directory (`cd Physion`), simply use

```bash
bash scripts/train_sgnn.sh [SCENARIO] [GPU_ID]
```
where [SCENARIO] can be selected from 3-wheel Cubli, 2-wheel Cubli AGVcar and so on, and [GPU_ID] is the id of the gpu available, such as 0.

*2. Evaluation.*

```bash
bash scripts/eval_sgnn.sh [SCENARIO] [EPOCH] [ITER] [SCENARIO] [GPU_ID]
```
where [EPOCH] and [ITER] are used to select the checkpoint, and setting them as zeros will automatically select the checkpoint with best validation loss.

