Unsupervised object-centric representation (OCR) learning has recently drawn attention as a new paradigm of visual representation. Despite its potential as a pre-training technique for image-based reinforcement learning (RL) tasks in terms of sample efficiency, systematic generalization, and reasoning, the advantages of OCR in RL haven't been systematically examined.

Our work delves into the effectiveness of OCR pre-training for image-based RL tasks, driven by empirical experiments and analyses through a series of questions and answers. We introduced a simple object-centric visual RL benchmark for a comprehensive evaluation. Our findings contribute valuable insights into OCR pre-training for RL, outlining its benefits and potential limitations in particular scenarios.

Questions we aimed to answer include:

• "Does OCR pre-training improve performance on object-centric tasks?"

• "Can OCR pre-training help with out-of-distribution generalization?"

The research examines critical aspects of incorporating OCR pre-training in RL, such as performance in visually complex environments and the selection of an appropriate pooling layer for aggregating object representations.

Simple Object-Centric Visual RL Benchmark

Our benchmark is meticulously crafted to assess the proficiency and boundaries of various RL models. It focuses on essential capabilities, such as object recognition, interaction, and relational reasoning.

Our benchmark comprises of diverse tasks categorized into five major types:

1. Object Goal Task: Navigate to a designated object (blue box) in the environment.

2. Object Interaction Task: Push the blue box in the middle of the observation to another blue box at the left bottom of the scene.

3. Object Comparison Task: Find and go to the distinct object among the objects.

4. Property Comparison Task: Find and go to the object with the unique property.

5. Object Reach Task: Object Goal Tasks on visually more complex environments.

These tasks span a broad spectrum of difficulty and complexity, providing a comprehensive evaluation of the models against diverse challenges.

Experimental Results

We present our experimental results and analysis as a series of questions and answers, each probing a different aspect of OCR pre-training for RL.

Question 4: Does OCR pre-training work well in visually complex environments where segmentation is difficult?

The agent utilizing SLATE demonstrated superior sample efficiency and converged success rate compared to the other methods. Although this task does not explicitly require reasoning among the objects, it is still crucial for the agent to learn to avoid touching the distractor objects before the target object.

Question 5: Which OCR model is better for RL?

From the figure on Question 1, It is clear that the SLATE model performs the best in terms of overall performance on the tasks evaluated in this study.

Question 6: How does the choice of pooling layer affect task performance?

The use of the MLP pooling layer resulted in inferior performance on all tasks, with a complete failure to solve the interaction and comparison tasks.

Question 8: Does OCR pre-training work better than the baselines in environments with more objects? What happens if there are fewer objects?

The binding problem in neural networks refers to the difficulty in representing multiple objects as distinct entities when they are encoded into a single vector representation. OCRs, on the other hand, provide a scalable solution to this problem, as they can represent multiple objects independently. 

The results show that as the number of objects increases, the performance of both models decreases. However, the performance degradation of the CNN model is much greater than that of the SLATE model. The results show that with fewer objects in the environment, both VAE and SLATE performed better, with the difference being more pronounced for the VAE model.

Question 11: Do the standard metrics of evaluating OCR correlate with RL performance?

In this study, we aimed to investigate the correlation between standard OCR evaluation metrics, such as segmentation quality, reconstruction loss, and property prediction accuracy, and RL performance.

The results indicate that while the FG-ARI showed a negative correlation with RL performance across all tasks, the correlation between MSE, property prediction accuracy, and RL performance was inconsistent, with positive correlation observed for some tasks and negative correlation for others. From these results, we hypothesize that when performance in segmentation or reconstruction is good enough, such as in the case of SLATE, the correlation with RL performance is weaker.

Conclusion and Discussion

We found more specific conditions to satisfy the hypotheses through this empirical investigation. OCR pre-training does not always provide sample efficiency but is efficient for relational reasoning tasks. OCR pre-training is not always better than other methods for out-of-distribution tasks also, but it shows better generalization performance for unseen objects when compared with GT.

We chose those visually simple scenes to ensure the downstream RL performance is not affected by poor segmentation quality. It allows us to probe more specific aspects of the reinforcement learning task to assess where OCR pre-training is most beneficial. In order to investigate the case where segmentation quality is not perfect, we also ran experiments on the robotics Object Reaching Task, which we discuss in Question 4. Investigating OCR in more complex and realistic environments is a promising direction for future work, especially as unsupervised OCR models continue to improve.

Lastly, we hope our benchmark can help evaluate OCR models in the context of agent learning!

You can discover additional details about the experiments and other intriguing experimental findings in our paper. We hope you find them enjoyable!