This document will mainly explain the logics in the implementations, including the formats of input data, how to construct and train the network, and how to qualitatively and quantitatively test the trained network. Through this document, we hope readers can learn about how to use and modify this repo based on their purposes.

Further reading: Flexible Neural Representation for Physics Prediction, FleX ML Agents Simulation Environment

As the README in the repo already shows how to start the network training using scripts/train.py, which we hope is clear and intuitive by itself about the main training framework, we will go directly into the details of the important modules of this framework.

Formats of input data

The provided dataset describes two 4*4*4 rigid cubes colliding with each other on a static plane described by 5,000 particles. To build this dataset, we generated multiple scenarios, where at the beginning two cubes were randomly teleported on the plane and then pairs of forces were applied to random particles in two cubes to push them together. In order to completely characterize the motion dynamics within these scenarios, we recorded the 128 particle positions in each time point along with other information such as their masses, which cube they belong to, and forces applied to them.
Additionally, we precomputed collisions between pairs of paricles not belonging to the same cube and collisions between 128 dynamic particles and 5000 static particles for the plane, as these computations are time-consuming if done online. Moreover, due to some weird phenomena during scenario switching, we also provide signals about whether the scenario has just begun so that these bad frames can be excluded.

To be able to flexibly load the data, we split it into multiple parts and store these parts in different folders. Generally within all these folders, data is stored in the format of tfrecords, where each file includes records each of which represents information at one time step. To successfully combine the split data later, we store the data in the same time order in these folders. More specifically, the five sub-folders in the provided dataset folder new_tfdata or new_tfvaldata represent five parts of data, introduced as following:

Folders collision and static_collision

Folder collision includes pre-computed pair-wise collision relations between two cubes. And folder static_collision contains pre-computed collision relations between dynamic particles and static plane particles. To enable efficient data loading, we only keep the closest 1000 collision relations. Within each relation (i, j, d), i, j represents particle indexes and d represents the Euclidean distance between them. Moreover, when relation (i, j, d) is included, (j, i, d) will not be.

Folders is_moving and is_acting

Besides these two collision folders, data in folder is_moving represents whether the scenario has just begun and data in folder is_acting indicates whether forces are applied at that time point.

Folder full_particles

Finally, folder full_particles comprises important physical states. Due to historical reasons, we structured this information of each particle as a 22-dimension array fp_t, where each dimension encodes information such as position, mass, and force. To be more specific, fp_t[0:3] is x, y, z positions. fp_t[3] is the inverse of mass, so that we can represent static particles by setting fp_t[3] to be 0. fp_t[4:7] is delta position, meaning that fp_t[0:3] - fp_t[4:7] = fp_t-1[0:3]. Besides, fp_t[7:10] is x, y, z force applied to this particle and fp_t[14] is object index, where 0 means invalid particle which can be due to paddings happened when multiple datasets are combined. Additionally, we also put the future delta position in this array at fp_t[15:18], meaning that fp_t[0:3] + fp_t[15:18] = fp_t+1[0:3]. We only use this future delta position for loss computation.

Network construction and training

More updates will come in this section, although we hope function hrn in models/interaction_model.py is self-explaining.

Network validation and results visualization

More updates will come in this section, although we hope script test.py in folder scripts and example.ipynb in folder visualize are self-explaining.

To cite our paper please use:

@inproceedings{mrowca2018flexible,
  title={Flexible Neural Representation for Physics Prediction},
  author={Mrowca, Damian and Zhuang, Chengxu and Wang, Elias and Haber, Nick and Fei-Fei, Li and Tenenbaum, Joshua B and Yamins, Daniel LK},
  booktitle={Advances in Neural Information Processing Systems},
  year={2018}
}

Further reading: Flexible Neural Representation for Physics Prediction, FleX ML Agents Simulation Environment