VicK9 / locs

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

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

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:

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

Datasets

• 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 temporarily downloaded from https://drive.google.com/file/d/1oYzC-zCYAFSSxSDkrwADcH5eHsZXgMf9/view?usp=sharing. In the future, this link will be replaced with a permanent one, and the source code to generate the data will be added to the repository.

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