Object-Category Aware Reinforcement Learning
Anonymous Author(s)
AffiliationAddress email
Abstract
Object-oriented reinforcement learning (OORL) is a promising way to improve thesample efficiency and generalization ability over standard RL.
Recent works thattry to solve OORL tasks without additional feature engineering mainly focus onlearning the object representations and then solving tasks via reasoning based onthese object representations.
However, none of these works tries to explicitly modelthe inherent similarity between different object instances of the same category.
Objects of the same category should share similar functionalities; therefore, thecategory is the most critical property of an object.
Following this insight, we propose a novel framework named Object-Category Aware Reinforcement Learning(OCARL), which utilizes the category information of objects to facilitate bothperception and reasoning.
OCARL consists of three parts: (1) Category-Aware
Unsupervised Object Discovery (UOD), which discovers the objects as well as theircorresponding categories; (2) Object-Category
Aware Perception, which encodesthe category information and is also robust to the incompleteness of (1) at thesame time; (3) Object-Centric Modular Reasoning, which adopts multiple independent and object-category-specific networks when reasoning based on objects.
Our experiments show that OCARL can improve both the sample efficiency andgeneralization in the OORL domain.
Introduction
Reinforcement Learning (RL) has achieved impressive progress in recent years, such as results in
Atari [24] and Go [28] in which RL agents even perform better than human beings.
Despite itssuccesses, RL is also known to be of low sample efficiency [11] and poor generalization ability [4].
To deal with these limitations, object-oriented RL (OORL) [13; 9; 23; 18; 38] is a promising way.
Inspired by studies in cognitive science [25] that objects are the basic units of recognizing the world,
OORL focuses on learning the invariant of objects’ functionalities in different scenarios to achievebetter generalization ability.
In OORL the agent’s observation is a set of object representations and the task can be solved viareasoning based on these objects.
Most previous works [13; 18; 9] in OORL employ hand-craftedobject features given in advance, which require human expertise and yield in low generality, limitingthe applications of these methods.
Recent works [38; 37; 36] try to broaden the scope of OORL’sapplications by avoiding using additional feature engineering.
Their works can be divided into eithertop-down approaches [38] which purely rely on reward signals to ground representations to objects,or bottom-up approaches [37; 36] which utilize unsupervised object discovery (UOD) technologiesto provide structural object representations.
However, none of these works [38; 37; 36] tries to explicitly model the inherent similarity betweendifferent object instances of the same category when reasoning based on objects.
In most cases,objects of the same category should share the same functionalities, such as obeying similar behaviourpatterns or having the same effects when interacting with other objects.
Therefore, an intelligent agent should recognize an object’s category to be aware of its functionalities.
Since objects fromdifferent categories differ in their functionalities, the objects’ functionalities can also be separatelymodelled (according to their categories), leading to modularity which is beneficial for generalization[6; 8] in general.
Following the above insight, we propose a framework named O bject- C ategory
A ware R einforcement
L earning (OCARL), which utilizes the category information of objects to facilitate both perceptionand reasoning.
OCARL consists of three parts: (1) category-aware UOD, which can automaticallydiscover the objects as well as their corresponding categories via UOD and unsupervised clustering.(2)
O bject- C ategory
A ware P erception (OCAP), a perception module that takes the object categoryinformation obtained in (1) as an additional supervision signal.
OCAP can encode the categoryinformation into representations and is also robust to the incompleteness of UOD at the same time.
O bject- C entric M odular R easoning (OCMR), a reasoning module that adopts multiple independentand object-category-specific networks, each of which processes the object features of the samecorresponding category.
Such a modular mechanism can enhance further the generalization ability of
OCARL.
Our experiments show that OCARL can improve the sample efficiency on several tasks andalso enable the agent to generalize to unseen combinations of objects where other baselines fail.
Related Works
Unsupervised Object Discovery Unsupervised object discovery (UOD) methods try to automatically discover objects without additional supervision.
Roughly speaking, there are two maincategories in this area: spatial attention models and spatial mixture models.
Spatial attention modelssuch as SCALOR [17] and TBA [16] explicitly factorize the scene into a set of object properties suchas position, scale and presence.
Then these properties are utilized by a spatial transformer network toselect small patches in the original image, which constitute a set of object proposals.
These methodscan deal with a flexible number of objects and thus are prominent in the UOD tasks.
Spacial mixturemodels such as IODINE [12] and MONet [3] decompose scenes via clustering pixels that belong tothe same object.
However, they often assume a fixed maximum number of objects and thus cannotdeal with a large number of objects.
SPACE [21], a UOD method used in this paper, combines both the spatial attention model and spatialmixture model.
The spatial attention model extracts objects from the scene, whereas the spacialmixture model is responsible for decomposing the remaining background.
Such a combinationenables SPACE to distinguish salient objects from relative complex backgrounds.
Object-Oriented Reinforcement Learning Object-oriented reinforcement learning (OORL) hasbeen widely studied in RL community.
There are many different structural assumptions on theunderlying MDP in OORL, such as relational MDP [13], propositional object-oriented MDPs [9],and deictic object-oriented MDP [23].
Although these assumptions differ in detail, they all share acommon spirit that the state-space of MDPs can be represented in terms of objects.
Following theseassumptions, [5] proposes object-focused Q-learning to improve sample efficiency.
Schema Networks[18] learn an object-based casual model for regression planning.
[31] improve the generalizationability to novel objects by leveraging relation types between objects that are given in advance.
Although these works have been demonstrated to be useful in their corresponding domain, theyrequire explicit encodings of object representations, which limits their applications.
There are other lines of work in OORL such as COBRA [33], OODP [39], and OP3 [30] that try tosolve OORL tasks in an end-to-end fashion.
Most of these works fall into a model-based paradigmin which an object-centric model is trained and then utilized to plan.
Although these methods areeffective in their corresponding domains, in this paper, we restrict ourselves to the model-free setting,which can generally achieve better asymptotic performance than model-based methods [32; 14].
The most related works should be RRL [38] and SMORL [37].
RRL proposes an attention-basedneural network to introduce a relational inductive bias into RL agents
Method
OCARL consists of three parts: (1) category-aware unsupervised object discovery (category-aware
UOD) module, (2) object-category aware perception (OCAP), and (3) object-centric modular reasoning (OCMR) module.
The overall architecture is shown in Figure 1.
We will introduce (1)(2)(3) in
Section 3.1, Section 3.2, and Section 3.3 respectively.
3.1 Category-Aware Unsupervised Object Discovery
In this section, we explain how to automatically discover the objects as well as their correspondingcategories.
To put it simply, the objects are discovered by SPACE (a UOD method), and theircategories are then predicted via unsupervised clustering.
Unsupervised Object Discovery
In this work, SPACE [21] is utilized to discover objects from rawimages.
We first run a random policy to obtain enough data and then apply SPACE to this collecteddataset.
In SPACE, an image x is decomposed into two latent representations: foreground z fg and backgroundz bg .
For simplicity, we only introduce the inference model of the foreground latent z fg , which is aset of object representations.
We encourage the readers to refer to [21] for more details.
To obtainz fg , x is treated as if it were divided into H × W cells and each cell is tasked with modelling atmost one (nearby) object.
Therefore, z fg consists of a set of H × W object representations (i.e.z fg = { z fgij } Hi =1 Wj =1 ), each of which is a 3-tuple 1 z fgij =
( z presij , z whereij , z whatij ) .
z presij ∈ { 0 , 1 } is a1-D variable that indicates the presence of any object in cell ( i, j ) , z whereij encodes z fgij ’s location,and z whatij identifies z fgij itself.
The whole inference model of z fg can be written as:
[Equation]
where x ij is a small patch from the x that is obtained by the proposal bounding box identified byz whereij , which should contain exactly one object if z presij = 1 .
Unsupervised Clustering After applying SPACE on the randomly collected dataset, we can obtaina set of object representations:
D = { z presnij , z wherenij , z whatnij , x nij } Nn =1 Hi =
Wj =1 , where N is the sizeof the dataset, z presnij , z wherenij , z whatnij , x nij are defined in Eq.(1).
We first select object patches from
D that is of high object-presence probability by a threshold τ : D τ = { x nij ∈ D : p ( z presnij =1) > τ } .
Roughly speaking, each patch x nij ∈ D τ should contain exactly one object.
Therefore,we first project x nij ∈ D τ into a low-dimensional latent space R d by running IncrementalPCA[26] on D τ , and then adopt KMeans clustering with a given cluster number C upon the projectedlatent representations.
The object-category predictor q ( z catij | x ij ) =
Ψ KMeans ◦ Ψ IncrementalPCA is thecomposition of IncrementalPCA and KMeans:
[Equation]
where z catij ∈ R C is an one-hot latent which indicates the category of z ij .
For simplicity, we regard the background (i.e. z presij = 0 ) as special ’object’, thus z catij can be writtenas ˆz catij =
Ψ concatenate ([1 − z
presij
, z presij · z catij ]) .
Note that ˆ z catij is an one-hot latent ∈ R C +1 .
Although the object-category predictor Eq.(2) is very simple, we find it works well in our experiments.
In more complex environments, we can also rely on other advanced unsupervised clustering methodssuch as NVISA [20], CC [19], etc.
3.2 OCAP: Object-Category Aware Perception
The OCAP module is designed to robustly incorporate the object knowledge of the (category-aware)
UOD model defined in Section 3.1.
In a word, OCAP is a plain convolution encoder that accepts the category information from the UOD model as additional supervision signals.
Such a design makes
OCARL robust to the incompleteness of the UOD model: in extreme cases in which the UOD modelfails to discover any objects, OCAP degenerates into a plain convolution encoder which is still ableto extract useful information from the raw images with the help of reward signals.
Suppose the convolution encoder in OCAP maps the raw image observation X ∈ R 3 × H img × W img toa latent representation Z ∈ R F × H × W , where H, W are the same with those in the UOD model (see
Eq.(1)).
On the other hand, X is also fed into the UOD model to get the (extended) object categoryinformation ˆ z cat =
{ ˆ z catij } Hi =1 Wj =1 via Eq.(2).
The OCAP module forces the latent representation Z to encode the object category information,which is implemented by training a additional category predictor f cat : R F → Categorical ( C +1) .f
cat predicts ˆ z catij given Z : ,i,j which is a single channel in Z .
Therefore, the additional supervisionsignal that incorporates the object knowledge of the UOD model is given as:
[Equation]
L cat is used for training the convolution encoder in OCAP.
In practice, L cat can be used as an auxiliary loss to any RL algorithm:
[Equation]
where λ cat ∈ R + is a coefficient.
3.3 OCMR: Object-Centric Modular Reasoning
OCMR is a module that takes Z ∈ R F × H × W from the OCAP module as input and outputs a featurevector that summarizes Z .
The key design philosophy of OCMR is to adopt multiple independentand object-category-specific networks, each of which focuses on processing the object features ofthe same corresponding category.
Compared with using a universal category-agnostic network, theprocessing logic of each independent network in OCMR is much simpler and therefore easier tomaster, which in turn allows for improved generalization as we will show in Section 4.3.2.
Given Z ∈ R F × H × W from OCAP, we append the x-y coordinates (corresponding to the H × Wgrid) to each channel, embed the resulting tensor ∈ R ( F +2) × H × W back into R F × H × W , and finallyrearrange it into (cid:98) Z ∈ R HW × F .
(cid:98) Z is treated as HW objects with x-y coordinate informationencoded.
To model the relations between objects, we apply a self-attention [29] module on (cid:98) Z toquery information for each object from its related objects:
[Equation]
where W q , W k , W v ∈ R F × F are trainable parameters.
For example, whether a door can be openedis determined by the existence of the key; therefore, the door should query other objects to checkwhether it is openable.
After the attention module, (cid:98) Q and (cid:98) Z are concatenate together to get (cid:101)
Z =
Ψ concatenate ([ (cid:98) Q , (cid:98) Z ]) ∈
R HW × 2 F .
Each (cid:101)
Z i,j, : in (cid:101)
Z is then fed into different independent neural networks according to itscorresponding object category ˆ z catij .
In practice, this can be implemented as:
[Equation]
where f c : R 2 F → R F , ˆ z catij ; c ∈ { 0 , 1 } is the probability that (cid:101)
Z i,j, : belongs to c ’th category.
Eq.(5)can be computed in parallel, resulting a tensor Z out ∈ R HW × F .
Finally, the Z out is added to (cid:98) Z (i.e. a residual connection).
We perform a max operation alongthe H, W dimension, and using the resulting vector ∈ R F to predict the value function and actionprobabilities:
[Equation]
Experiment
4.1 Task Description
In this work, we consider two task domains:
Crafter and Hunter .
The observations onboth tasks are raw images of shape 64 × 64 .
Crafter [15] is a complex 2-D MineCraftstyle RL task, where complex behaviors are necessary for the agent’ survival.
The environmentis procedurally generated by arranging variousresources, terrain types, and objects (18 in total).
The agent is rewarded if it can craft newitems and accomplish achievements (22 possibleachievements in total).
[Table]
The ’Score’ is a metric proposed in [15] takes the difficultiy of each achievement thus more
Hunter is much simpler than Crafter .
It is also procedurally generated but only contains 4 typesof objects: Hunter , Cow , Zombie and Wall 2 .
The agent can control Hunter and can get positivereward (=1) if Hunter catches a Cow or kill a Zombie .
The agent will be given a high reward (=
5)if it can accomplish this task by catching/killing all Cow s and Zombie s.
However, once Hunter iscaught by a Zombie , the agent will receive a negative reward (=-1), and the episode ends.
In a word,the agent in Hunter should master two different behavior patterns: chase & catch the Cow , and avoid& shoot at Zombie .
Although Hunter is simple, it is a typical OORL task.
We can derive different environment instancesfrom Hunter by setting the number of Zombie s and Cow s. We use Hunter-ZmCn to denote an environment that spawns ( m Zombie s + n Cow s) at the beginning of each episode, and Hunter-ZmCn/ZnCman environment that spawns ( m Zombie s + n Cow s) or ( n Zombie s + m Cow s).
4.2 Evaluation
In this section, we evaluate OCARL’s effect both on sample efficiency and generalization ability.
We consider the following baselines: PPO [27], RRL [38], and SMORL [37].
PPO is a generalpurpose on-policy RL algorithm with a plain neural network (i.e. Convolution + MLP).
RRLproposes an attention-based neural network to introduce relational inductive biases and iteratedrelational reasoning into RL agent, which aims to solve OORL tasks without any other externalsupervision signals.
SMORL utilizes the UOD model to decompose the observation into a set ofobject representations which are directly used as the agent’s new observation.
All algorithms (RRL,
SMORL, and OCARL) are (re-)implemented upon PPO to ensure a fair comparison.
For moreimplementation details, please refer to the Appendix.
4.2.1 Sample Efficiency
The results on Hunter domain are shown in Figure 3.
We can conclude three facts from this figure: (1)
OCARL performs better and is more stable than other baselines on all tasks.
(2) RRL is competitivewith OCARL on Hunter-Z4C4 but significantly worse on other simper environments.
This is because
RRL relies solely on the reward signals to ground the visual images into objects; therefore, it requiresenough (rewarded) interactions between objects, which simpler environments such as Hunter-Z1C1fail to supply.
(3) SMORL performs better than other baselines (PPO, RRL) on simple environmentssuch as Hunter-Z1C0/Z0C1 , Hunter-Z1C1 and Hunter-Z4C0/Z0C4 , but unable to make progress in Hunter-Z4C4 which consists of more objects.
This is because SMORL adopts a simple reasoningmodule that may not be well-suited to multi-object reasoning.
The results on Crafter domain are shown in Table 1 and Figure 4.
Crafter is a complex environment in that it features sparse (rewarded) interactions between objects and large object categorynumbers.
Besides, some objects are omitted by the UOD model (see Appendix) because they almostdo not appear in the training dataset due to the insufficient exploration of the random policy.
Due tothese features, both RRL and SMORL fail to provide noteworthy results (with scores of 4.22 and 3.94, respectively).
On the other hand, OCARL is able to present more meaningful behaviours (witha score of 12.31), such as collecting coal, defeating zombies, making wood pickaxes and so on, asshown in Figure 4.
The learning curves and detailed achievement success rates are provided in the
Appendix.
4.2.2 Generalization
To evaluate the generalization ability, we train agents on Hunter-Z1C0/Z0C1 , Hunter-Z1C1and Hunter-Z4C0/Z0C4 separately, and then observe their test performance on Hunter-Z4C4 .
In Hunter-Z1C1 , the agent needs to generalize from a few objects to more objects.
In
Hunter-Z4C0/Z0C4 , the agent never observes the coincidence of Cow and Zombie .
The generalization from Hunter-Z1C0/Z0C1 to Hunter-Z4C4 is the most difficult.
Such a paradigm actually follows the out-of-distribution (OOD) setting, in that the object combination ( 4 Zombie s + 4 Cow s)is never seen in the training environments.
Therefore, although such generalization is possible forhumans, it is much harder for RL algorithms.
As shown in Figure 5, OCARL is the only algorithm that make significant progress on the testenvironment.
Although other baselines can improve their training performance (Figure 3), they areunable to generalize to the test environment.
The generalization performance of OCARL is mostimpressive on Hunter-Z1C1 , because agent trained on Hunter-Z1C1 even achieves slightly betterthan that on Hunter-Z4C4 (i.e. directly trained on the test environment, see the last figure in Figure3.).
4.3 Ablation study
4.3.1 Robustness to the Incompleteness of Object Discovery Methods
OCARL relies on UOD methods to discover different objects from given observations, which arenot guaranteed to discover all objects, especially in complex environments [34].
Therefore, one maynaturally ask to what extent OCARL’s performance is affected by such incompleteness.
To answer this question, we randomly mask out each object (by setting z presij to 0 ) discovered by the
UOD model with a given probability p ∈ [0 , 1] , and evaluate OCARL’s performance after training 5Menvironment steps on Hunter .
Such random confusion of the UOD model will affect both the OCAPand OCMR modules of OCARL.
Note that OCARL( p = 0 ) is OCARL itself, and OCARL( p = 1 )is exactly RRL.
The results are shown in Figure 6.
Although increasing p does hinder OCARL’sperformance, OCARL can still benefit from the imperfect UOD model in that its performance isstill better than RRL in general.
Such robustness should be attributed to the fact that OCARL doesnot directly use the output of the UOD model as observation but instead treats it as an additionalsupervision signal that help the encoder to capture object-category information, which gives OCARLa second chance to find other useful information from the raw image.
4.3.2 OCMR Improves Generalization
One key advantage of OCARL is that it can generalize to unseen combinations of objects by leveragingthe invariant relations between objects learnt during training.
We argue that such an advantage isestablished via OCMR, which adopts object-category-specific networks to deal with different objects.
To confirm our argument, we consider a variant of OCARL (i.e. OCARL w/o OCMR) in which the
OCMR module is removed and replaced by an RRL-like reasoning module which adopts a universalnetwork (instead of multiple object-category-specific networks as OCARL does).
The comparisonbetween OCARL w/ and w/o OCMR is shown in Figure 7.
OCARL w/o OCMR (almost completely)fails to generalize from Hunter-Z1C0/Z0C1 , Hunter-Z4C0/Z0C4 to Hunter-Z4C4 , while OCARLw/ OCMR does.
In Figure 7, we also consider a variant of OCARL (i.e. OCARL w/o OCAP) inwhich the OCAP module is disabled by setting λ cat = 0 in Eq.(3).
In this paradigm, only OCMR
The mean returns over 12 seeds are plotted with a 95% confidence interval.
Without the help of
OCMR, the generalization ability of OCARL will decline significantly.
On the other hand, OCMR alone can bring benefits without the help of OCAP (although not as consistently as OCMR+OCAP).
[Table]
is in effect. As shown in Figure 7, OCMR alone can still bring benefits without the help of OCAP(although not as consistently as OCMR+OCAP ).
According to these results, we can conclude that
OCMR is crucial in OCARL’s success.
4.3.3 The Modularity of OCMR
OCMR adopts different modules to deal with object features from different categories.
In this section,we will show that each module in OCMR actually encodes an object-category-specific behaviourpattern.
As stated in Section 4.1, the agent should master two useful behaviour pattern in Hunter : (1) chase& catch the Cow , and (2) avoid & shoot at Zombie .
The existence of these patterns can be observedby evaluating the agent in Hunter-Z0C4 and Hunter-Z4C0 respectively, because agents that masterpattern (1) (pattern (2)) should achieve better performance on Hunter-Z0C4 ( Hunter-Z4C0 ).
In
Table 2, we first train agent on Hunter-Z4C4 and then disable the module f Z (correspondingto Zombie ) or f C (corresponding to Cow ) by replacing f C or f E with f bg (corresponding to the’background’) in Eq.(5).
The results shows that disabling f Z does not affect the performance on
Hunter-Z0C4 ( 8 . 29 → 8 . 13 ), which means that f Z is not relative to behaviour pattern (2).
Instead,f Z is corresponding to behaviour pattern (1) because we can observe a significant performancedecline ( 8 .
→ 1 . 17 ) on Hunter-Z4C0 when we disable f Z .
Similar phenomenon can also beobserved on f C .
Conclusion
In this paper, we propose OCARL, which utilize the category information of objects to facilitateboth perception.
For the perception, we propose the OCAP, which enables the encoder to distinguishthe categories of objects.
For the reasoning, we propose the OCMR, which adopts object-categoryspecific networks to deal with different objects.
Our experiments are carried out on Crafter and
Hunter .
Experiments show that OCARL outperforms other baselines both on sample efficiency andgeneralization ability.
We also perform several ablation studies to show that (1) OCARL is robust tothe incompleteness of UOD methods, (2) OCMR is critical to improving generalization, and (3) eachmodule in OCMR actually encodes an object-category-specific behaviour pattern in Hunter .
Checklist
The checklist follows the references.
Please read the checklist guidelines carefully for information onhow to answer these questions.
For each question, change the default [TODO] to [Yes] , [No] , or[N/A] .
You are strongly encouraged to include a justification to your answer , either by referencingthe appropriate section of your paper or providing a brief inline description.
For example:
• Did you include the license to the code and datasets?
[Yes] See Section ?? .
• Did you include the license to the code and datasets?
[No] The code and the data areproprietary.•
Did you include the license to the code and datasets?
[N/A]
Please do not modify the questions and only use the provided macros for your answers.
Note that the
Checklist section does not count towards the page limit.
In your paper, please delete this instructionsblock and only keep the Checklist section heading above along with the questions/answers below.
For all authors...
(a) Do the main claims made in the abstract and introduction accurately reflect the paper’scontributions and scope?
[Yes](b) Did you describe the limitations of your work?
[No](c) Did you discuss any potential negative societal impacts of your work?
[No](d) Have you read the ethics review guidelines and ensured that your paper conforms tothem?
[Yes]2.
If you are including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results?
[N/A](b)
Did you include complete proofs of all theoretical results?
[N/A]3.
If you ran experiments...
(a) Did you include the code, data, and instructions needed to reproduce the main experimental results (either in the supplemental material or as a URL)?
[Yes](b) Did you specify all the training details (e.g., data splits, hyperparameters, how theywere chosen)?
[Yes] see Appendix.(c)
Did you report error bars (e.g., with respect to the random seed after running experiments multiple times)?
[Yes]
(d) Did you include the total amount of compute and the type of resources used (e.g., typeof GPUs, internal cluster, or cloud provider)?
[No]
If you are using existing assets (e.g., code, data, models) or curating/releasing new assets...
(a)
If your work uses existing assets, did you cite the creators?
[Yes]
(b) Did you mention the license of the assets?
[Yes] Partially.
(c) Did you include any new assets either in the supplemental material or as a URL?
[Yes]
(d) Did you discuss whether and how consent was obtained from people whose data you’re
using/curating?
[No]
(e) Did you discuss whether the data you are using/curating contains personally identifiable
information or offensive content?
[No]
If you used crowdsourcing or conducted research with human subjects...
(a) Did you include the full text of instructions given to participants and screenshots, ifapplicable?
[N/A](b)
Did you describe any potential participant risks, with links to Institutional Review
Board (IRB) approvals, if applicable?
[N/A]
(c) Did you include the estimated hourly wage paid to participants and the total amountspent on participant compensation?
[N/A]