Codes for reproducing main results in paper "Flexible Neural Representation for Physics Prediction".
We have tested this repo under Ubuntu 16.04 with tensorflow version 1.9.0.
For now, we provide one dataset including two rigid cubes each of which contains 64 particles hitting each other on a static plane described by 5000 particles.
You can download the dataset through this link.
For those who downloaded the dataset before Nov. 29th, please download it again as we updated the file to include pretrained model and validation files.
After downloading, untar this file.
Please check README.md in data
folder for more explanations about this dataset.
Go to folder scripts
. And run the following command:
sh train_physics.sh --gpu your_gpu_number --dataset path_to_your_dataset
By default, the models and log will be saved to ~/.model_cache/physics_pred/
.
Besides, multi-gpu training is not supported for now.
You can restore your training by setting restore_path
parameter.
During training, the reported losses including loss
, preserve_distance_loss
, un_velocity_loss
, and velocity_loss
should usually keep decreasing.
At the end of the training, velocity_loss
should be around 0.0006
.
Additionally, we provide two example bash scripts for quantatively and qualitatively verifying the trained model. Besides, we also provide a pretrained model trained using the training bash script in the dataset.
Now if you want to run a qualitative test, you can go to folder scripts
. And run the following command:
sh test_physics_qual.sh --gpu your_gpu_number --dataset path_to_your_dataset --restore_path path_to_your_dataset/pretrained_model/checkpoint-384000
By default, this will generate two pickle files named as true_results_physics_pred_better_8.pkl
and results_physics_pred_better_8.pkl
in your home directory.
You can change saving directly by setting SAVE_DIR
parameter. Please check examples of visualizing these results in visualize
folder.
Similarly for a quantitative test, you can just replace test_physics_qual.sh
using test_physics_quan.sh
.
Then you will see the metrics also reported in our paper as outputs of this command.
And the test result will be stored as quant_results_physics_pred_better_8.pkl
in your home directory by default.