B INDING A CTIONS TO O BJECTS IN W ORLD M ODELS
A BSTRACT
We study the problem of binding actions to objects in object-factored world models using action-attention mechanisms.
We propose two attention mechanisms for binding actions to objects, soft attention and hard attention, which we evaluate in the context of structured world models for five environments.
Our experiments show that hard attention helps contrastively-trained structured world models to learn to separate individual objects in an object-based grid-world environment.
Further, we show that soft attention increases performance of factored world models trained on a robotic manipulation task.
The learned action attention weights can be used to interpret the factored world model as the attention focuses on the manipulated object in the environment.
Source code: https://github.com/ondrejba/action attention iclr 22.
I NTRODUCTION
Reinforcement learning agents often operate in domains that contain multiple objects, such as robotic manipulation (e.g. Janner et al.
; Li et al.
or game playing (e g. Bellemare et al.
; Guss et al.
An action performed by an agent in an object-based environment is unlikely to affect all objects at once.
A household robot might pick up a single object or push several objects but it will not move all objects in a room at once.
In Atari Pong, the agent only controls the paddle while the other two objects–the ball and the other player–are not directly influenced by actions.
The sparsity of interactions between actions and objects is a useful inductive bias for an agent’s world model.
We implement a learned binding of an action to objects, which is constrained to be sparse.
The sparsity is enforced by an attention mechanism over objects, with individual weights being constrained to sum to one using the softmax operator.
Prior work considered structured action spaces in world models for grid-worlds with objects (Kipf et al., 2020), a Tetris-like block stacking task (Janner et al., 2019; Veerapaneni et al., 2019)
and in model-free learning for 2D construction from blocks (Bapst et al., 2019).
Unlike our work, the binding of an action to a particular object, or an edge between objects in a graph, was hand-coded, without any learning.
Goyal et al.
(2021) proposed an attention mechanism for binding production rules (Lovett & R. Anderson, 2005), which include information about action selection, to slots in a structured model.
We propose two different attention mechanisms and test them in object-based grid worlds, two
Atari games and a simulated robotic manipulation task.
Firstly, we work with contrastively-learned structured world models (C-SWMs, Kipf et al.
that jointly learn to factor states and predict the dynamics of the environment.
This class of models often requires a strong bias in order to learn to properly distinguish objects without explicit supervision.
To this end, we encode the bias that an action affects exactly one object using a categorical random variable.
We call this approach hard attention , drawing on computer vision literature (Sharma et al., 2015).
Secondly, we test our idea in a realistic robotic manipulation domain, where the individual states are already factored into object and our task is to learn factored world models (FWMs, Biza et al.
Here, the robot is tasked with building towers from blocks; all blocks in a tower might potentially be affected by a robot’s action.
Hence, the categorical action binding bias is too restrictive.
Instead, we propose a soft attention module, which transforms and weights actions for each object based on their predicted impact on said object.
It is based on single-head self-attention (Vaswani et al., 2017).
We limit our exploration of action binding to C-SWMs and FWMs; see Battaglia et al.
; Greff et al.
for a review of other structured model architectures.
Our results indicate two positive outcomes and one negative outcome.
When each object can be individually affected by an action, hard action attention proves an excellent inductive bias for learning to factor states.
Conversely, when some objects are not directly affected by actions, neither soft nor hard action attention makes a difference.
Finally, we show that soft attention can lead to large improvements in a pick-and-place robotic manipulation task.
In summary, our main contributions are: #1: We propose two novel attention modules over actions, soft and hard attention, that improve the ability of structured models to factor states and to predict transition dynamics.
#2: We identify the phenomenon wherein a structured world model captures highly correlated information in all its available slots.
This happens in Atari Pong and Space Invaders, where only the player-controlled paddle or spaceship is directly affected by actions.
P RELIMINARIES : S TRUCTURED AND F ACTORED W ORLD M ODELS
World models are trained to make predictions about the state of an environment s 0 after applying a sequence of actions a 1 , a 2 , ... a T .
It is commonly assumed that the state is represented as an image and the transition dynamics are deterministic.
We make no assumptions about the action space.
World models used in this paper learn a latent space without image reconstruction.
That is, the model encodes s 0 into a latent encoding z 0 and predicts the future z 1: T based on a 1: T .
Structured world models encode an image s into a latent state composed of K slots z 1: K , each represented as a feature vector.
The model learns to assign objects in s to individual slots (i.e. to factor the scene) without any direct supervision.
We use the contrastively-trained structured world model (C-SWM, Kipf et al.
(2020))–it learns structured embeddings using a convolutional encoder
E ϕ and predicts the effects of actions using a graph neural network (GNN) T
θ (Gori et al., 2005).
The GNN uses a f node and an edge network f edge to compute the next latent state of the k th object:
[Equation]
The model is trained by contrastive learning using the following loss:
[Equation]
The negative example ¯ z is an encoding of a state randomly sampled form the training set and γ is a hyper-parameter.
Intuitively, the negative term in the loss states that a pair of randomly sampled states should be at least γ apart in the embedding space.
This prevents the model from converging to a trivial solution.
We use factored world models to refer to world models that operate over factored states.
That is, they learn to model the transition dynamics of individual factors without having to learn to discover factors.
Biza et al.
(2022) proposed a factored world model for pick-and-place robotic manipulation tasks.
The model receives a factored state s 1: K composed of images of the K object present in a scene.
Each state factor is individually encoded into a latent factor: z 1: K = ⟨ f enc ( s 1 ) , f enc ( s 2 ) , ..., f enc ( s K ) ⟩ .
The transition model and the training loss follow C-SWM with the exception of two design decisions in the transition model.
(1) several layers of GNNs are used and skip connections are added in between layers and (2) the action is provided not only to the node network f node but also to the edge network
f edge .
We experiment with different numbers of GNN layers, but omit the edge actions.
B INDING A CTIONS TO O BJECTS
In this section, we describe two attention mechanisms that bind actions to object slots.
In both cases, we start with a factored latent state z 1: K , which is an output of an encoder of either C-SWM or
FWM.
We want to compare the individual factors to the action in order to determine which slots are affected.
First, we transform individual factors into key vectors, k = ⟨ j k ( z 1 ) , j k ( z 2 ) , ..., j k ( z K ) ⟩ , and the action into a query vector q = j q ( a ) .
Both j k and j q are Multi-Layer Perceptrons.
The keys and queries are used to compute attention weights
α over object slots 1 to K for action a :
[Equation]
It is important to note that in both the hard and soft attention modules described below, the attention weights end up transforming the action that is given to the node network of the GNN transition model.
For the k th object slot, instead of computing f node ( z k , a, ... )
we will use f node ( z k , a ′ k , ... ) , where a ′ k is an action constructed specifically for the k th object slot.
Hard attention: We define a categorical variable M that represents the membership of an action to a particular factor:
[Equation]
This random variable only plays a role when predicting the transitions and computing the positive term of the contrastive loss (the negative term only depends on the current encoded state).
During training, we compute the expectation of M via summation over all of its possible assignments (same as the number of object slots / factors).
During inference, we take the assignment with the highest probability.
The random variable is used as follows:
[Equation]
with the pad ( a t , m ) operator assigning action a t to the m th slot and zeros to all other slots.
Finally, the positive term of the contrastive loss function is computed as follows:
[Equation]
Soft attention: Unlike hard attention, soft attention weights individual actions provided to slots by attention weights without defining a probabilistic model.
We follow the single-head self-attention module from the Transformer (Vaswani et al., 2017).
The action a is first transformed into a value vector v =
j v ( a ) .
Then, the value vector is multiplied by the previously computed attention weights for each factor: a ′ =
⟨ α 1 v, α 2 v, ..., α K v ⟩ .
We then pass the (newly) factored action a ′ into the transition model.
Note the individual factors in the action are all equal up to a multiplicative factor.
Table 2 in the Appendix visualizes α in a robotic pick-and-place task.
E XPERIMENTS
We aim to answer the following questions:
Does hard attention help C-SWM factor environments where each object can be independently affected by an action ( 2D Shapes and 3D Cubes )?
Does the same apply when some objects are only indirectly (or not at all) affected by actions (Atari Pong and Space Invaders)?
Does soft attention increase the performance of factored world models in robotic manipulation?
E NVIRONMENTS
We use two grid-world ( 2D Shapes and 3D Cubes ) and two Atari environments (Atari Pong and Space Inveders) from Kipf et al.
(2020) and a robotic manipulation environment from Biza et al.
We make changes to 2D Shapes and 3D Cubes in order to better study the effect of action attention.
We describe each environment below.
• 2D Shapes and 3D Cubes: In both environments, five objects of different shapes and colors are placed in a 5 × 5 grid- world.
Each object can be moved independently and the objects are not allowed to pass through each other.
Crucially, Kipf et al.
(2020) represented actions as a pair of object index and direction of movement.
As the action space was factored by design, CSWM could use this information to make state factorization easier.
We change the action space to represent object position and direction of movement.
Hence, the model does not know the index of the object being moved, making state factorization more challenging.
[Table]
• Atari Pong and Space Invaders : C-SWM is trained and tested on modeling short action sequences (10 timesteps) in these two games.
We use the approach from Biza et al.
(2021), where a trained agent acts for a random number of steps to reach an interesting starting state from which a random policy is rolled out.
The collected datasets contain only the random rollouts so that the agent experiences both good and bad actions.
• Robotic Pick-and-Place: A simulated UR5 robotic arm manipulates six cubes.
The action space is the continuous space of ( x, y ) coordinates, where a pre-programmed pick or place action is executed.
The environment is observed by two side-viewing RGB cameras and additional two RGB cameras capture the content of the robot’s hand.
These observations are post-processed to extract individual objects and their bounding boxes to capture each object independently (i
e. a factored state space).
The training task is to make a tower of four cubes; the model is then expected to transfer to four testing tasks–wall, stairs, two towers of three and three towers of two–without additional training.
4. S HAPES , C UBES AND A TARI G AMES
Setup: The C-SWMs are evaluated on their ability to predict future latent states in an evaluation dataset using ranking metrics (Hits@1 and MRR, see Appendix A.1).
These metrics check if the predicted latent state is closer to the encoding of the true next state compared to a set of negative states.
Intuitively, we both test the ability of the model to predict the future and the ability of the encoder to distinguish between states.
Result:
Ranking scores are reported in Table 1.
We would like to highlight the performance of hard action attention in 2D Shapes and 3D Cubes .
In these environments, baseline C-SWM fails because the action space is not factored per individual object.
In contrast, C-SWM with hard attention encodes the correct inductive bias and learns to assign individual objects to separate slots in most cases.
We do not see the same increase in performance in Atari games: baseline C-SWM, soft attention and hard attention achieve comparable results.
To explain this result, we introduce a new metric to measure the correlation between individual slots.
If all slots capture the same information given a single input state, the correlation coefficient is going to be 1 (please see Appendix A.1).
We see that successful models in 2D Shapes and 3D Cubes have a low slot correlation, as each individual object should be represented by a different slot.
Conversely, all models in Atari (as well
[Figure]
as failed models in 2D Shapes and 3D Cubes ) have high correlation between slots.
We show one example of this phenomenon in Figure 2 in the Appendix where the latent space of all slots is nearly identical (the first slot seems to be flipped, but follows the same pattern).
The correlation metric results suggest that C-SWM has no need for more than one slot in Atari games.
This is supported by the original results from Kipf et al.
We hypothesize that since the action often affects only one object (the paddle in Pong and the spaceship in Space Invaders), there is not enough information to guide C-SWM to factor objects not directly affected by the action.
4. R OBOTIC P ICK - AND -P LACE
Setup: We replicate the experiment from Biza et al.
(2022), where a UR5 robotic arm picks and places six cubes.
A dataset of expert demonstrations is available for the training task and the model is evaluated on its ability to predict the block positions and rank action sequences in expert demonstrations of the zero-shot transfer tasks.
The block position prediction metric measures if groundtruth information about the positions of all blocks can be decoded from the learned latent space.
The action ranking metric measures if the model can distinguish action sequences that successfully solve a task from perturbed action sequences that do not.
See Appendix A.1 for additional information.
Result: Factored world models with soft action attention outperform the baseline models when the number of GNN layers is set to one or two (Figure 1).
As the number of layers increases to three and four, the performance of the two models converges.
The attention weights predicted by Soft
Attention for a model with one graph neural network layer are shown in Table 2 in the Appendix.
The model correctly identifies which object is being manipulated.
Additionally, the model pays attention to where a particular cube is being placed.
When building a tower, the cubes below the placed object also receive the information about the place action with a low (but non-zero) weight.
We believe this behavior emerged because placement of a cube can often move the cubes below.
Furthermore, we hypothesize that the diminishing benefit of soft attention as we add more GNN layers can be explained by the increasing capacity of the model.
A single GNN layer can benefit from the pre-computed information about which object will be affected.
On the other hand, a stack of four GNN layers receives the information about the action at each layer, making it easier to distinguish affected and unaffected objects within the transition model.
C ONCLUSION
We propose two attention mechanisms for binding actions to objects, soft attention and hard attention, and validate them on the task of learning structured world models in five environments.
As one of our reviewers pointed out, ”there seems to be a dilemma between soft and hard attention – the soft attention is more expressive, however, it also under-performs hard attention in some experiments.”
We view this trade-off from the perspective of expressiveness versus bias.
Hard attention provides a strong bias (that all actions move exactly one object), which allows C-SWM to converge to the right solution in the grid-world environments.
Conversely, soft attention is more flexible in providing information about the action to several objects.
This construction fits our simulated robotic manipulation task, where several cubes might move at once when they are being stacked. A CKNOWLEDGEMENTS This work is supported in part by NSF 1724257, NSF 1724191, NSF 1763878, NSF 1750649, NSF 1835309, NSF 2107256 and NASA 80NSSC19K1474.
A A PPENDIX A. E VALUATION M ETRICS
C-SWM Hits@1 and Mean Reciprocal Rank: We follow the evaluation setup from Kipf et al.
An evaluation dataset of episodes of length 10 is collected.
The agent’s encoder is used to encode all states in the evaluation dataset into latent states.
Then the agent predict the next latent state t
= 1 , t = 5 or t = 10 steps into the future given a sequences of 1, 5 or 10 actions respectively.
The predicted latent state is then compared to the actual encoded state at timestep t (the positive example) and to all states at timestep t from other episodes (the negative example).
The
Hits@1 metric simply counts the fraction of times the predicted latent state was closer to the positive example than to all other negative examples.
The Mean Reciprocal Rank (MRR) metric computes
[Equation]
where D is the evaluation dataset and rank n is the index of the positive examples in the sorted list of distances between the predicted latent state and the set of the positive and all negative examples.
The original evaluation code does not always handle duplicate states in different episodes well.
It is assumed that two identical images will be encoded into the same latent state (i.e. all float32 values will be exactly equal), but this does not happen in practice even when pytorch is set to a deterministic mode.
Duplicate states are infrequent in 2D Shapes and 3D Cubes , but fairly common in Atari.
In the Atari experiments, we check for duplicate states and remove negative states that are identical to a particular positive state.
Slot Correlation: we compute the absolute value of the Pearson correlation coefficient between every pair of object slots.
This value is averaged over all pairs of object slots and over all states in the evaluation set.
The resulting value is bounded between zero and one, where zero means no correlation between slots and one means that all slots capture the same information.
FWM Block Position RMSE and Action Sequence Hits@1: We use the same metrics as Biza et al.
To the block position prediction, we train a Multi-Layer Perceptron to predict the spatial ( x, y, z ) position of each block from the latent space of FWMs.
Only the MLP is trained, all other weights are frozen.
Then, we compute the root mean squared error between the actual and predicted block positions.
The dataset used in this evaluation is the dataset of expert demonstrations for the four zero-shot transfer tasks.
In action sequence ranking, we generate one correct and ten incorrect trajectories for each zeroshot transfer task.
We also provide the model with an example of a goal state.
The model encodes the goal state and predicts the resulting latent state of all eleven action sequences.
If the predicted result for the correct trajectory is the closest to the encoded goal state, the prediction of the model is considered correct.
Hits@1 once again measures the fraction of correct predictions.
[Table]