# A General and Efficient SE(3)-Equivariant Graph Framework: Encoding Symmetries with Complete Differential Invariants and Frames
## Environment requirements
Install DGL: https://www.dgl.ai/pages/start.html
 
## N-Body Simulation
### Generate datasets
For example, to generate a dataset with 5 balls, follow these steps:
```
cd n_body_system
cd dataset
```
and
```
python generate_dataset.py --n_balls 5
```
### Run experiments
*CDIF model*
```
python -u main_n_body.py --exp_name cdif_1 --model cdif --lr 0.001 --num_agents 5 --nf 32
```

*EGNN model*
```
python -u main_n_body.py --exp_name egnn_vel_1 --model egnn_vel --lr 5e-4 --num_agents 5 --nf 32
```

*EGNN model + Complete Invariants*
```
python -u main_n_body.py --exp_name egnn_vel_cis_1 --model egnn_vel_cis --lr 0.001 --num_agents 5 --nf 32
```

*GMN model*
```
python -u main_n_body.py --exp_name gmn_1 --model gmn --lr 5e-3 --num_agents 5 --nf 32
```

*GMN model + Complete Invariants*
```
python -u main_n_body.py --exp_name gmn_cis_1 --model gmn_cis --lr 0.001 --num_agents 5 --nf 32
```

*ClofNet model*
```
python -u main_n_body.py --exp_name clof_vel_1 --model clof_vel --lr 0.001 --num_agents 5 --nf 32
```

*ClofNet model + Complete Invariants*
```
python -u main_n_body.py --exp_name clof_vel_cis_1 --model clof_vel_cis --lr 0.001 --num_agents 5 --nf 32
```

*Radial Field*
```
python -u main_n_body.py --exp_name rf_vel_1 --model rf_vel --lr 2e-4 --num_agents 5 --nf 32
```

*TFN Model*
```
python -u main_n_body.py --exp_name tfn_1 --model tfn --lr 1e-3 --degree 2 --nf 32 --num_agents 5
```

*SE3 Transformer*
```
python -u main_n_body.py --exp_name se3_1 --model se3_transformer --lr 5e-3 --div 1 --degree 3 --nf 32 --num_agents 5
```
## Formation Consensus 
### Generate datasets
This dataset is partitioned into four distinct difficulty levels: **easy**, **medium**, and **difficult**. You have the flexibility to specify the number of agents according to your specific requirements.

For example, to generate a dataset at the **difficult** level with 5 agents, follow these steps:
```
cd con_formation
```
and
```
python generate_cf.py --level difficult --n_balls 5
```
### Run experiments
Assume that the level of datasets is **difficult**.

*CDIF model*
```
python -u main_cf.py --exp_name cdif_difficult --model cdif --level difficult --lr 0.001 --outf con_formation/logs --num_agents 5
```

*ClofNet model*
```
python -u main_cf.py --exp_name clof_vel_difficult --model clof_vel --level difficult --lr 0.001 --outf con_formation/logs --num_agents 5
```

*EGNN model*
```
python -u main_cf.py --exp_name egnn_vel_difficult --model egnn_vel --level difficult --lr 0.001 --outf con_formation/logs --num_agents 5
```

*GMN model*
```
python -u main_cf.py --exp_name gmn_difficult --model gmn --level difficult --lr 0.0008 --outf con_formation/logs --num_agents 5
```

*Radial Field*
```
python -u main_cf.py --exp_name rf_vel_difficult --model rf_vel --level difficult --lr 2e-4 --outf con_formation/logs --num_agents 5
```

*TFN Model*
```
python -u main_cf.py --exp_name tfn_difficult --model tfn --level difficult --lr 1e-3 --degree 2 --outf con_formation/logs --num_agents 5
```

*SE3 Transformer*
```
python -u main_cf.py --exp_name se3_difficult --model se3_transformer --level difficult --div 1 --lr 5e-3 --degree 3 --outf con_formation/logs --num_agents 5
```
### Experiments on invariant complexity
Exploring the relationship between the complexity and quantity of invariants and the expressiveness of networks.
Case: 1 to 6 correspond to 6 different models, which vary in the quantity and complexity of their invariants. 
```
python -u main_cf.py --exp_name cdif_invar_difficult --model cdif_diff_invar --level difficult --lr 0.001 --outf con_formation/logs --num_agents 5 --case 1
```

## MD17 
The MD17 dataset can be downloaded from [MD17](http://www.sgdml.org/#datasets). Put the downloaded file in "md17/dataset".
### Run experiments
*CDIF model*
```
python -u main_md17.py --exp_name cdif_1 --model cdif --lr 0.001 --nf 64
```

*EGNN model*
```
python -u main_md17.py --exp_name egnn_vel_1 --model egnn_vel --lr 5e-4  --nf 64
```

*GMN model*
```
python -u main_md17.py --exp_name gmn_1 --model gmn --lr 5e-3 --nf 64
```

*ClofNet model*
```
python -u main_md17.py --exp_name clof_vel_1 --model clof_vel --lr 0.001 --nf 64
```

*Radial Field*
```
python -u main_md17.py --exp_name rf_vel_1 --model rf_vel --lr 2e-4 --nf 64
```

*TFN Model*
```
python -u main_md17.py --exp_name tfn_1 --model tfn --lr 1e-3 --degree 2 --nf 64 --num_agents (the corresponding number of atoms)
```
## CMU Motion Capture
### Run experiments
*CDIF model*
```
python -u main_motion.py --exp_name cdif_1 --model cdif --lr 0.001 --nf 64 --delta_frame 30
```

*EGNN model*
```
python -u main_motion.py --exp_name egnn_vel_1 --model egnn_vel --lr 5e-4  --nf 64 --delta_frame 30
```

*GMN model*
```
python -u main_motion.py --exp_name gmn_1 --model gmn --lr 5e-3 --nf 64 --delta_frame 30
```

*ClofNet model*
```
python -u main_motion.py --exp_name clof_vel_1 --model clof_vel --lr 0.001 --nf 64 --delta_frame 30
```

*Radial Field*
```
python -u main_motion.py --exp_name rf_vel_1 --model rf_vel --lr 2e-4 --nf 64 --delta_frame 30
```

*TFN Model*
```
python -u main_motion.py --exp_name tfn_1 --model tfn --lr 1e-3 --degree 2 --nf 64 --delta_frame 30
```
