Official source code for

Roto-translated Local Coordinate Frames for Interacting Dynamical Systems
Miltiadis Kofinas, Naveen Shankar Nagaraja, Efstratios Gavves
NeurIPS 2021
https://arxiv.org/abs/2110.14961

Local frames construction demo:

LoCS GNN demo:

We propose local coordinate frames for all objects in the geometric graph of an interacting dynamical system. Each local coordinate frame is centered at the position of the target object and rotated such that it matches the object's orientation. Our method, termed LoCS (Local Coordinate FrameS), can be added to any graph network that takes as inputs nodes positioned in space, and evolving through time.

The source code is written in Pytorch and uses PyTorch Geometric.

Create a Conda environment using environment.yml.

conda env create -f environment.yml

Then, install the project by running:

pip install -e .

Alternatively, you can use the install.sh script to install the environment.

./install.sh

The script installs all dependencies one by one. Please mind that the script might install different versions of some secondary dependencies.

This codebase is based on the code from the papers:

• Neural Relational Inference for Interacting Systems, https://github.com/ethanfetaya/NRI
• Dynamic Neural Relational Inference, https://github.com/cgraber/cvpr_dNRI

• Synth: this code includes the synthetic data, as well as code used to generate it.
• Motion Capture: the datasets can be downloaded from http://mocap.cs.cmu.edu/search.php?subjectnumber=35. For subject 35, you need trials 1-16 and 28-34.
• InD: Data must be requested from here: https://www.ind-dataset.com/
• Charged Particles: The dataset can be downloaded from https://doi.org/10.5281/zenodo.7500066. Eventually, the source code to generate the data will be added to the repository.

Scripts train models can be found in the scripts directory. For all datasets, the script names have the following structure: run_{dataset}_{model}.sh.

You can train and evaluate LoCS in the synthetic dataset by running the following command from the repository root directory:

./run_synth_locs.sh

Similarly, you can train and evaluate NRI, dNRI, and EGNN, by changing the model name.

You can train and evaluate LoCS in the charged particles dataset by running the following command from the repository root directory:

./run_charged_locs.sh

Similarly, you can train and evaluate NRI, dNRI, and EGNN, by changing the model name.

You can train and evaluate LoCS in the motion capture dataset by running the following command from the repository root directory:

./run_motion_35_locs.sh

Similarly, you can train and evaluate NRI, dNRI, and EGNN, by changing the model name.

You can train and evaluate LoCS in inD by running the following command from the repository root directory:

./run_ind_locs.sh

Similarly, you can train and evaluate GRU, dNRI, and EGNN, by changing the model name.

If you use this code or this model in your work, please cite our work:

@inproceedings{kofinas2021roto,
  title={Roto-translated Local Coordinate Frames For Interacting Dynamical Systems},
  author={Kofinas, Miltiadis and Nagaraja, Naveen Shankar and Gavves, Efstratios},
  booktitle={Thirty-Fifth Conference on Neural Information Processing Systems (NeurIPS)},
  year={2021}
}