Awesome

PrAE

This repo contains code for our CVPR 2021 paper.

Abstract Spatial-Temporal Reasoning via Probabilistic Abduction and Execution
Chi Zhang*, Baoxiong Jia*, Song-Chun Zhu, Yixin Zhu
Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition (CVPR), 2021
(* indicates equal contribution.)

Spatial-temporal reasoning is a challenging task in Artificial Intelligence (AI) due to its demanding but unique nature: a theoretic requirement on representing and reasoning based on spatial-temporal knowledge in mind, and an applied requirement on a high-level cognitive system capable of navigating and acting in space and time. Recent works have focused on an abstract reasoning task of this kind -- Raven’s Progressive Matrices (RPM). Despite the encouraging progress on RPM that achieves human-level performance in terms of accuracy, modern approaches have neither a treatment of human-like reasoning on generalization, nor a potential to generate answers. To fill in this gap, we propose a neuro-symbolic Probabilistic Abduction and Execution (PrAE) learner; central to the PrAE learner is the process of probabilistic abduction and execution on a probabilistic scene representation, akin to the mental manipulation of objects. Specifically, we disentangle perception and reasoning from a monolithic model. The neural visual perception frontend predicts objects' attributes, later aggregated by a scene inference engine to produce a probabilistic scene representation. In the symbolic logical reasoning backend, the PrAE learner uses the representation to abduce the hidden rules. An answer is predicted by executing the rules on the probabilistic representation. The entire system is trained end-to-end in an analysis-by-synthesis manner without any visual attribute annotations. Extensive experiments demonstrate that the PrAE learner improves cross-configuration generalization and is capable of rendering an answer, in contrast to prior works that merely make a categorical choice from candidates.

Performance

The following table shows the performance of various methods on the RAVEN and I-RAVEN datasets. For details, please check our paper.

Performance on RAVEN / I-RAVEN:

Dependencies

Important

• Python 3.8
• PyTorch
• CUDA and cuDNN expected

See requirements.txt for a full list of packages required.

Usage

To train the PrAE learner, one needs to first extract rule annotations for the training configuration. We provide a simple script in src/auxiliary for doing this. Properly set path in the main() function, and your dataset folder will be populated with rule annotations in npz files.

To train the PrAE learner, run

python src/main.py train --dataset <path to dataset>

The default hyper-parameters should work. However, you can check main.py for a full list of arguments you can adjust.

In the codebase, window sliding and image preprocessing are delegated to the dataset loader and the code only supports training on configurations with a single component.

One thing we notice after code cleaning is that curriculum learning is not necessary, but in the manuscript we keep our original discovery.

To test on a new configuration, run

python src/main.py test --dataset <path to dataset> --config <new config> --model-path <path to a trained model>

Testing on 3x3Grid could potentially raise a CUDA-out-of-memory error. Try running on CPU then.

Citation

If you find the paper and/or the code helpful, please cite us.

@inproceedings{zhang2021abstract,
    title={Abstract Spatial-Temporal Reasoning via Probabilistic Abduction and Execution},
    author={Zhang, Chi and Jia, Baoxiong and Zhu, Song-Chun and Zhu, Yixin},
    booktitle={Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition (CVPR)},
    year={2021}
}

Acknowledgement

We'd like to express our gratitude towards all the colleagues and anonymous reviewers for helping us improve the paper. The project is impossible to finish without the following open-source implementations.

• RAVEN
• WReN
• CoPINet
• LEN
• MXGNet
• I-RAVEN & SRAN