/locs

Official source code for "Roto-translated Local Coordinate Frames For Interacting Dynamical Systems". In NeurIPS 2021

Primary LanguagePythonMIT LicenseMIT

Roto-translated Local Coordinate Frames for Interacting Dynamical Systems

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

LoCS

TL;DR

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.

Install

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

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 the dependencies.

Attribution

This codebase is based on the code from the papers:

Datasets

Experiments

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

Synthetic Dataset

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.

Charged Particles

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.

Motion #35

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.

inD

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.

Cite

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}
}