Title: IMAGINE THAT! LEVERAGING EMERGENT AFFORDANCES FOR TOOL SYNTHESIS IN REACHING TASKS Anonymous

Abstract: In this paper we investigate an artificial agent's ability to perform task-focused tool synthesis via imagination. Our motivation is to explore the richness of information captured by the latent space of an object-centric generative model -and how to exploit it. In particular, our approach employs activation maximisation of a task-based performance predictor to optimise the latent variable of a structured latent-space model in order to generate tool geometries appropriate for the task at hand. We evaluate our model using a novel dataset of synthetic reaching tasks inspired by the cognitive sciences and behavioural ecology. In doing so we examine the model's ability to imagine tools for increasingly complex scenario types, beyond those seen during training. Our experiments demonstrate that the synthesis process modifies emergent, task-relevant object affordances in a targeted and deliberate way: the agents often specifically modify aspects of the tools which relate to meaningful (yet implicitly learned) concepts such as a tool's length, width and configuration. Our results therefore suggest, that task relevant object affordances are implicitly encoded as directions in a structured latent space shaped by experience.

1. INTRODUCTION:

Deep generative models are gaining in popularity for unsupervised representation learning. In particular, recent models like MONet (Burgess et al., 2019) have been proposed to decompose scenes into objectcentric latent representations (cf. Greff et al., 2019; Engelcke et al., 2019) . The notion of such an objectcentric latent representation, trained from examples in an unsupervised way, holds a tantalising prospect: as generative models naturally capture factors of variation, could they also be used to expose these factors such that they can be modified in a task-driven way? We posit that a task-driven traversal of a structured latent space leads to affordances emerging naturally as directions in this space. This is in stark contrast to more common approaches to affordance learning where it is commonly achieved via direct supervision or implicitly via imitation (e.g. Tikhanoff et al., 2013; Myers et al., 2015; Liu et al., 2018; Grabner et al., 2011; Do et al., 2018) .

The setting we choose for our investigation is that of tool synthesis for reaching tasks as commonly investigated in the cognitive sciences. The ability to make and use tools has been a cornerstone of humanity's cognitive, technological and cultural development. It has long been hailed as a defining characteristic of intelligence (Ambrose, 2001) . Consequently, the ability of artificial agents to select and use appropriate tools has received significant attention in, for example, the robotics community (e.g. Sinapov & Stoytchev, 2007; Tikhanoff et al., 2013; Toussaint et al., 2018; Xie et al., 2019) . A very limited number of works such as Wicaksono (2017) also consider tool making and typically deploy strong inductive biases such as simulations in which the agent adjusts numerical parameters given a set of design constraints. However, while much stands to be gained by endowing robots with an aptitude for tool use, more recent evidence in the cognitive sciences suggests that habitual tool use cannot in and off itself be taken as a sign of intelligence (e.g. Ambrose, 2001; Emery & Clayton, 2009) . Some animals, which are recognised as prodigious users of tools, do not show any sign of understanding either cause or effect or, indeed, the difference between appropriate and inappropriate tools (e.g. Ambrose, 2001) . Instead, relatively recent results suggest that tool selection and manufacture -especially once demonstrated -is a significantly easier task than tool innovation: the step, prior to manufacture, of imagining the type of tool suitable for a particular task (Beck et al., 2011) . We posit that it is tool innovation which ultimately unlocks the full potential of tool manufacture and, thus, creative tool use.

In order to demonstrate that a task-aware, object-centric latent space encodes useful affordance information we require a mechanism to train such a model as well as to purposefully explore the space. To this end we propose an architecture in which a task-based performance predictor (a classifier) operates on the latent space of an object-centric generative model (see Figure 1 ). During training the classifier is used to provide an auxiliary objective, aiding in shaping the latent space. Importantly, however, during test time the performance predictor is used to guide exploration of the latent space via activation maximisation (Erhan et al., 2009; Zeiler & Fergus, 2014; Simonyan et al., 2014) , thus explicitly exploiting the structure of the space. While our desire to affect factors of influence is similar in spirit to the notion of disentanglement, it contrasts significantly with approaches such as β-VAE (Higgins et al., 2017) , where the factors if influence are effectively encouraged to be axis aligned. Our approach instead relies on a high-level auxiliary loss to discover the direction in latent space to explore.

Our experimental design is inspired by a long history in the cognitive sciences and behavioural ecology of investigating reaching tasks which require tool use (e.g. Kohler, 1925; van Leeuwen et al., 1994; Chappell & Kacelnik, 2002; 2004; Emery & Clayton, 2009; Beck et al., 2011) . The results demonstrate that artificial agents are indeed able not only to select -but to imagine -an appropriate tool for a variety of reaching tasks. The dataset consisting of 52,000 synthetic example reaching tasks will be made available to the community.

To the best of our knowledge, ours is the first work to demonstrate an artificial agent's ability to imagine, or synthesise, images of tools appropriate for a given task via optimisation in a structured latent embedding. Similarly, while activation maximisation has been used to visualise modified input images before (e.g. Mordvintsev et al., 2015) , we believe ours to be the first to effect deliberate manipulation of factors of influence by chaining the outcome of a task predictor to the latent space, and then decoding the latent representation back into an image. Beyond the application of tool synthesis, in demonstrating that object affordances can be viewed as directions in a structured latent space as well as by providing a novel architecture adept at deliberately manipulating interpretable factors of influence we believe our work to provide novel perspectives on affordance learning and disentanglement.

2. RELATED WORK:

The idea of an affordance, which describes a potential action to be performed on an object (e.g. a doorknob affords being turned), goes back to Gibson (1977) . Together, robotics and computer vision boast a rich literature on learning task-affordances (e.g. Tikhanoff et al., 2013; Mar et al., 2015; Stoytchev, 2005) , which includes work on learning multiple attention-like affordance masks (provided labelled data) for common objects such as kitchen utensils (Do et al., 2018; Myers et al., 2015) . To learn affordances without supervision, another research thread has employed simulations in virtual environments. Grabner et al. (2011) is an example in which a simulated humanoid was re-positioned, via trial-and-error, to find all of the surfaces in a 3D scene in which it is possible to sit. Other studies have successfully learned affordances through imitation (e.g. Liu et al., 2018) .

Closely related to the the affordance literature is a mature body of research dealing with almost all aspects of robot tool use. Tasks addressed include reaching (Jamone et al., 2015) , grasping (Takahashi et al., 2017) , pushing (Stoytchev, 2005) and hammering (Fang et al., 2018) . This research commonly depends on a more traditional pipeline approach, consisting of tool-recognition (Tikhanoff et al., 2013; Fang et al., 2018) , toolselection (Xie et al., 2019; Saito et al., 2018) , planning (Toussaint et al., 2018) and execution (Tikhanoff et al., 2013; Fang et al., 2018; Stoytchev, 2005; Jamone et al., 2015) . The toolkit is an image of stick-like tools (here in yellow, light green, and pink). The goal of the task is to determine if any tool can reach the dark green target whilst touching the red boundary and avoiding any blue obstacles. When no object in the toolkit satisfies these geometric constraints, our model (TasMON) is able to take a tool from the toolkit and imagine suitable modifications until it does satisfy these constraints (right). Here we see TasMON imagine the yellow stick transforming into a hook with a thin enough handle to both pass through the gap between the blue obstacles and reach the dark green target, all whilst still crossing over the red boundary.

Recently, the advent of deep generative models has led to an active research area using world models (Ha & Schmidhuber, 2018; Lesort et al., 2018; Nair et al., 2018) , in which an artificial agent can train itself on tasks using a kind of imagination. Our work is directly inspired by this, though we approach imagination modelling in a different way. In particular, we propose an architecture in which a generative model for unsupervised scene decomposition (Burgess et al., 2019) is paired with the use of activation maximisation (e.g. Erhan et al., 2009; Zeiler & Fergus, 2014; Simonyan et al., 2014) of an auxiliary, task-driven loss-signal back into the generative model's latent representations to drive the imagination process.

3. DATASET:

To investigate tool imagination, we designed a set of simulated reaching tasks with clear and controllable factors of influence. Each task image is comprised of a green target button, a red line delineating the workspace area, and, optionally, a set of blue obstacles. We also vary the goal location and the sizes and positions of the obstacles. For each task image, we provide a second image depicting a set of 'tools' -straight sticks, L-shaped hooks, and J-shaped claws -with varying dimensions, shapes, colours, and poses. Given a pair of images (i.e. the task image and the toolkit image), the goal is to select a tool for a given task scene that can reach the target (the green dot) whilst avoiding obstacles (blue areas) and remaining on the exterior of the workspace (i.e. behind the red line). Depending on the task image, the applicability of a tool is determined by different subsets of its attributes. For example, if the target button is unobstructed, then any tool of sufficient length will satisfy the constraints (regardless of its width or shape). However, when the target is hidden behind a corner, or only accessible through a narrow gap, an appropriate tool also needs to feature a long-enough hook, or a thin-enough handle, respectively. By design, the colour of a tool does not influence its applicability and is introduced as a nuisance factor. As depicted in Figure 2 , we have designed eight scenario types to study these factors of influence in isolation and in combination. Importantly, the scenario types are designed to enable an investigation into our approach's performance on in-sample tasks, interpola- We choose these scenarios to set constraints on the dimensions of an appropriate tool; for example, in scenario C (third column), a tool needs to be thin enough to reach the goal. Other attributes which affect applicability are tool length and type. In the top row, we overlay tools in yellow, which can press the buttons without colliding with the obstacles or leaving the movement space, as indicated by the green tick. The bottom row, marked with red crosses, shows tasks being attempted by orange tools, which cannot succeed without violating the constraints. Tools are rendered here for visualisation and are not present in the task images.

tion tasks and an extrapolation task. In-sample scenarios differ from interpolation scenarios in that the tasks are different but can be solved with tools (sticks or hooks) that were seen during training. The extrapolation scenario consists of a novel scenario but almost always also requires the use of a new tool: a J-shaped claw.

For each task, toolkit images are constructed by scattering tools from a catalogue with randomised colours and poses in a separate visible workspace. A geometric applicability check, described in Appendix B, verifies whether or not any of the tools can reach the target, while satisfying the task constraints. We consider a task-toolkit pair feasible if the toolkit image includes at least one tool that can achieve the task under the specified constraints.

More formally, our dataset can be written as a set of N triplets:

{(I n G , I n T , ρ n )} N n=1

, where each example features a task image I G , a toolkit image I T , and a binary label ρ indicating the feasibility of reaching the target with one of the given tools (see Figure 3 ). We also save instance segmentation masks M n k for the toolkits, as well as the applicability label ρ n k for each tool, indexed by k. In all our experiments, we restrict the training input to the sparse high-level triplets and use the additional ground-truth labels for evaluation purposes only. For additional details on the dataset, see Appendix B.

4. METHOD:

We frame the challenge of tool imagination as an optimisation problem in a structured latent space obtained using generative models. The optimisation is driven by a high-level, task-specific performance predictor, which assesses whether a goal position is reachable given a particular tool. To reflect real-world complexity, the system must choose a tool from a set of potentially unsuitable tools (Figure 1 ) to prime the optimisation. In effect, we ask how an unsuitable tool might be improved to succeed on a task. This motivates an objectcentric, generative model with a latent space suited to task success classification as detailed in Figure 3 . 

4.1. OBJECT-CENTRIC TOOL REPRESENTATION:

Given that our tools are presented in toolkit images I T (see Figure 3 ), it is necessary for all subsequent processing to perform a decomposition of I T , from pixels into discrete objects. To achieve this decomposition of images into their component objects, we employ an unsupervised architecture called MONet (Burgess et al., 2019) . MONet consists of two parts: an attention network and a component VAE (Kingma & Welling, 2014; Rezende et al., 2014) . The attention network recurrently proposes K regions to represent individual parts of an image. These components hopefully correspond to individual objects (plus the scene background). Each component is represented as a VAE latent code and used to explain a certain part of the input image by reconstructing it. MONet therefore serves two purposes at once: its encoder ψ represents an image as a set of object-centric latent vectors, and its decoder ψ can be used to render any such latent vector into an individual object. We train MONet using the image-reconstruction loss L M ON et described in Burgess et al. (2019) . We denote the set of latent variables encoding the toolkit computed by the MONet encoder, ψ, as

ψ(I T ) = {z T,k } | k ∈ {1..K}.

(1)

After encoding the toolkit image, we need to represent the task image I G in an abstract latent space. The task encoder, φ, consists of a stack of convolutional layers followed by two dense layers.1 φ takes a concatenation of the task image I G and a normalised meshgrid of the image's (x, y) coordinates as input and maps this into the task embedding z G . Without loss of generality the (x, y) meshgrid can be generated for any input and provides φ with a sense of relative scale via coordinate querying (see Appendix D for further details).

In order to interpret the task success signal, we need to select the component that dominates the feasibility of the toolkit. This is done via a soft-attention approach which passes each tool representation z T,k concatenated with the task embedding z G as z cat,k through a three-layer MLP f to compute an attention logit (cf. Figure 3 ). A softmax activation is then applied to the logits to normalise the attention score such that ,j ) .

α k = e f (z cat,k ) Σ K j=1 e f (zcat

(2)

The attention over z cat,k is indicative of the influence of the k-th tool on global task success. We use the approximated soft attention to compute a toolkit context vector z containing information about all the tools and their respective contributions to the task solution, where

EQUATION

4.2. TOOL IMAGINATION:

Task-driven learning The context vector z contains both task-relevant information (e.g. tool length, width, and shape), as well as task-irrelevant information such as tool location and colour. In order to perform tool imagination, the sub-manifold of the latent space that corresponds to the task-relevant features has to be found. This is achieved by adding a three-layer MLP as a classifier σ. The classifier σ takes a concatenation h cat of the task embedding and the context vector as input and predicts the softmax over the binary task success. The classifier learns to identify the task-relevant sub-manifold of the latent space by using the sparse success signal ρ and optimising the binary-cross entropy loss, such that

EQUATION

where ρ ∈ {0, 1} is a binary signal indicating whether or not the toolkit contains a feasible tool to solve the task. The whole system is trained end-to-end with a loss given by

EQUATION

The hyper-parameter λ denotes a weight factor balancing the task and MONet losses. Note that the gradient from the task classifier σ propagates through both the task encoder φ and the toolkit encoder ψ, and therefore helps to shape the latent representations of the toolkit with respect to the requirements for task success.

Tool imagination Once trained, TasMON can synthesise new tools by traversing the latent manifold of individual tools in directions that maximise classification success given a toolkit image (Figure 3 ). To do this, we first concatenate each individual toolkit component z k with the task embedding z G and select the tool embedding vector z T,k that corresponds to the highest tool utility σ(z cat,k ) (i.e. the tool most likely to succeed). This warm-starts the imagination process. We denote the chosen tool latent as z im and its concatenation with the task embedding vector as z cat,im . We then use activation maximisation (Erhan et al., 2009; Zeiler & Fergus, 2014; Simonyan et al., 2014) to optimise the tool encoding z im with regard to the L task of the success estimation σ (z cat,im ) and a feasibility target ρ s = 1, such that

EQUATION

where η denotes the learning rate for the update. Finally, we apply this gradient update for S steps or until the success estimation σ (z cat,im ) reaches a threshold γ, and use ψ (z im ) to generate the imagined tool represented by z im .

5. EXPERIMENTS:

Our aim in this section is to demonstrate TasMON's ability to imagine a suitable tool given a task image and an image of a putative set of tools. Concretely, we examine the model's performance when imagination requires interpolating between properties of known tool types as well as when it requires innovating new tool types that have not previously been seen during training. To provide insight into the key aspects of TasMON's operation, our main evaluation in Table 1 focuses on imagination success, while we present performance on key sub-tasks such as tool utility, decomposition, and selection as supporting work in Appendix A.

Tool Utility (TU) We compute tool utility as the average precision aggregated over classifier scores σ for all tools when compared to their ground truth feasibilities. Details are provided in Appendix A.

Tool Imagination (TI):

We evaluate whether our model can generate tools to achieve tasks. For each instance, the target signal ρ is set to success such that the latent vector of the tool is then modified via backpropagation using a learning rate of 0.01 for 5, 000 steps or until σ(z cat,im ) reaches 0.997. The imagined tool is then rendered via the latent space decoder ψ and tested using the same geometric applicability test as described in Appendix B. We report TI as the percentage of imagined tools that successfully pass this test.

5.1. QUANTITATIVE RESULTS:

In order to investigate the efficacy of TasMON under different operating conditions, we present models trained with two different training regimes. TasMON is trained in curriculum learning manner: we first pre-train the MONet weights on toolkit images and then train the performance predictor with the MONet backbone jointly, i.e. the gradient from the predictor is allowed to back-propagate into the MONet. This curriculum allows MONet to learn object-centric decompositions before the performance predictor overfits the training data. To investigate the impact of jointly training the latent representation with the performance predictor we report results on an ablation model FroMON in which we keep the pre-trained MONet weights frozen during predictor training. Both TasMON and FroMON are trained on in-sample scenario types only. We select trained checkpoints that have the best performance prediction (TU) on a validation split. In addition to comparing the imagination outcomes for TasMON and FroMON, we also include a simple Random Walk (RW) baseline, where in place of taking steps in the direction of the steepest gradient, we move in a random direction in the latent space for 5000 steps. Specifically, the latent vector of the selected tool is updated by a sample drawn from an isotropic Gaussian with mean 0 and the absolute value of the groundtruth gradient derived by back-propagating from the predictor as the variance. For 250 infeasible instances per scenario type, we warmstart each imagination attempt with the same tool across RW, TasMON and FroMON to enable a like-for-like comparison. We also test to see how the models modify feasible tools; for 250 instances per scenario type, imagination is warmstarted with a feasible tool to observe how it is modified. The reader is referred to Table 2 and Appendix A for more quantitative results examining the model's performance in related sub-tasks like toolkit decomposition and tool selection.

TasMON successfully imagines applicable tools in 46.0% of in-sample test cases, and in 11.3% of tool interpolation test cases (Table 1 ). A significant drop in TI performance is incurred for the interpolation scenarios (E,F,G) in which the model encounters distinctly novel scenarios. The model fails in the extrapolation scenario where it needs to imagine a previously unseen tool type. TasMON outperforms FroMON in tool utility prediction (Table 2 ) and tool imagination in most tasks, suggesting that, although the predictor is powerful enough to guide the imagination through an already existing underling structure of toolkit-representations, a task-aware latent space can still provide benefits. As expected, the random walk fails in most scenarios except A and B. This is because there are no obstacles in scenario A, so any blob large enough to reach the target will suffice. We conjecture that a similar situation occurs in scenario B when the obstacle is small, which imposes only a weak constraint on the required tool. Interestingly, the imagination could also destroy some of the already feasible tools as showed in the TI-Feasible column in Table 1 . Upon examination, we find that this is due to the combined effect of underconfident behaviour from the performance predictor, reconstruction errors, and rendering choices. We refer readers to Appendix E for further details about rendering.

5.2. QUALITATIVE RESULTS:

Qualitative examples of the tool imagination process are provided in Figure 4 as well as in Figure 6 in the supplementary material. A striking feature of the optimisation process is the fact the performance predictor often drives the tools to evolve in a deliberate and interpretable way by modifying aspects such as length, width and configuration. This suggests that these properties are encoded as directions in the structured latent space learnt by TasMON and deliberately traversed via a high-level task objective in the form of the performance predictor. Also, since the tools are modified in a smooth manner, we hypothesise that tools are embedded in a continuous manifold of changing length, width and configuration. Moreover, instead of always turning a stick to a hook, the model can adjust the tool shape between these two configurations according to the tasks as depicted in block (c) of Figure 4 , which indicates that the model learns a non-biased interpretation of the tasks.

A ADDITIONAL QUANTITATIVE RESULTS:

In this appendix, we present performance on key sub-tasks useful for understanding model behaviour.

Tool Utility (TU) The performance predictor σ is responsible for estimating the probability of the toolkit containing an applicable tool for the given task based on the encoding of the task image concatenated with toolkit context vector. However, a toolkit image need not depict more than a single tool, and so σ can provide a measure of tool utility of an individual, possibly imagined, tool. We use intersection-over-union (IOU) as a similarity measure to match tool decompositions to ground truth tool masks according to the Kuhn-Munkres algorithm (Kuhn, 1955; Munkres, 1957) . All matches below the conservative threshold of IOU < 0.8 are discarded as they correspond to spurious or poor quality decompositions. Matched ground-truth masks are assigned their corresponding tool decomposition's tool utility probability from the classifier. We report the aggregated tool utility as average precision. A tool utility probability of 0 is assigned to any unmatched ground truth masks. Since we report average precision, this rubric can only deflate our scores by increasing false negatives at the cost of true positives, leading to conservative estimates of performance.

Toolkit Decomposition (TD) We use the Adjusted Rand Index (ARI) (Rand, 1971; Hubert & Arabie, 1985) to evaluate the quality of the instance segmentation, as in Greff et al. (2019) . An ARI score treats the pixelto-instance-assignment as a clustering problem and computes the similarity between the predicted clustering and ground truth. ARI scores range from 0 (complete misalignment) to 1 (identical clustering).

Tool Selection (TS) The TasMON model performs tool selection as described in Section 4.2. In feasible scenarios, we select the tool decomposition with the highest tool utility, and compare it with the ground truth mask of the applicable tool. We consider tool selection to be successful if the predicted mask has an IOU overlap with the ground-truth mask of at least 0.8, which we take to be a conservative threshold.

Table 2 : Results table, comparing TasMON with a frozen MONet+Task ablation, i.e. FroMON, across all tasks. Other abbreviations: TU = tool utility, aggregated as the average-precision score over all ground truth tools; TD = toolkit decomposition, measured by the Adjusted Rand Index (Rand, 1971; Hubert & Arabie, 1985) ; TS = tool-selection success in feasible scenarios, % of times pixel-wise IOU of most-likely tool and the ground-truth feasible tool was 0.8 or more; TI = unconditional tool-imagination success for each model, % of all test scenarios; The best model for a task and metric are shown in bold, except for TD where ARI scores were similar for different models. All metric scores are shown as percentages for consistency. Like-for-like tool-imagination results are shown in We see that both models demonstrate a strong ability to gauge tool utility (Table 2 , TU column) in in-sample scenarios. However, prediction performance decreases as scenario types grow more distinct from the training data. Table 2 suggest a strong correlation between successful tool imagination (TI) and effective and robust tool utility prediction (TU). This is unsurprising as the performance predictor constitutes the main driver for the optimisation. Finally, we note that task decomposition (TD) is commensurate across all models and scenario types. This is an intuitive result as both models are trained with an explicit decomposition objective, which is independent of scenario type.

B DATASET CONSTRUCTION:

Figure 5 : More illustrative examples of tasks (cf. Figure 2 ). We procedurally generated 196,390 unique instances across eight scenario types for our reaching experiments.

Our synthetic dataset consists of pairs of tasks and toolkits rendered from top-down views as 64 × 64 pixel RGB images. We create the dataset in multiple sampling steps ensuring a balance between scenarios and tool configurations. For each scenario, we generate 5,000 tasks. Each is paired with 50 tools from a tool catalogue containing 1 million unique tool geometries, evenly balanced between sticks, hooks and claw tools. Then we perform the tool applicability check for each of the task-tool pairs and reject tasks which had fewer than five or more than 40 applicable tools. The applicability of a tool to a task is determined by sampling 200 interior points of the tool polygon, overlaying the sampled point with the target and rotating the tool polygon and a vertically mirrored copy of it by 360 degrees. If any such pose of the tool satisfies all constraints (i.e. touching the space behind the red line while not colliding with any obstacle), we consider the tool applicable for the given task. In this way, we construct up to five feasible and five infeasible toolkits for each task, by sampling one to three tools per toolkit and rendering them with randomised poses and colours. We constrain the feasible toolkits to contain exactly one applicable tool. All toolkits are constrained to have only sticks and hooks as tools, with the exception of the H tasks, which require claws in order to be feasible.

The dataset contains a total of 52,000 scenarios. Table 3 shows a breakdown of each by scenario type and split. Each split has an equal number of feasible and infeasible instances, for each scenario type. As described in the main text, our model (TasMON) comprises two main parts (as depicted in Figure 3 ). On the one hand, we used a task-based classifier (Tas), while on the other we used a MONet as a scene decomposition model (MON). We briefly describe each component in turn.

C.1.1 MONET:

As a scene decomposition model, we faithfully re-implement MONet (Burgess et al., 2019) with the minor modification of using leaky ReLUs (Maas et al., 2013) as activation functions in place of ReLUs.

C.1.2 TASK-BASED CLASSIFIER:

The task-based classifier consists of three sub-parts: a task encoder, an attention module, and a classifier.

The task encoder stacks three convolutional blocks followed by a single convolutional layer and a two-layer MLP. First, the three consecutive convolutional blocks make use of 16, 32, 64 output channels respectively. Each convolutional block contains two consecutive sub-convolutional blocks, each with a kernel size of 3 and padding of 1 (and with stride of 2 and stride of 1, respectively). Second, the additional convolutional layer has stride 2, padding 1, and 128 output channels. All convolutional layers use a leaky ReLU nonlinearity. Finally, the MLP takes a flattend output of the convolutional layers of dimension 2048 (4 × 4 × 128) and passes it to a hidden layer with 128 neurons using a leaky ReLU as the activation function. This is then passed through another layer, giving a 16 dimensional output for the task-based classifier.

The attention module is a three layer MLP with input dimension 32, and successive hidden layers of 128 and 16 neurons respectively. Each of these layers use leaky ReLU activation functions. To predict the logit of the attention probability, the attention module has an output dimension of 1.

The structure of the classifier part of the task-classifier model mirrors the attention module. In place of the output layer, it consists of a three layer MLP with the same dimensions. Given that the classifier outputs logits for a binary classifier, the output dimension is 2.

C.2 TRAINING DETAILS:

All experiments are performed in PyTorch, using the ADAM optimiser with a learning rate of 0.0001 and a batch size of 32. We vary the weight factor λ between 0.1 to 4, and we select TasMON to be the best performing model, at λ = 1. MONet's hyperparameters tuned to our dataset are given in Table 4 . The CNN is local and its parameters are shared across spatial locations; however, the location of an object in an image is global information. In order to infer the global location of an object, the CNN must work together with the permutation-variant dense-layers. Assume that in a CNN-Dense network structure the output shape of the feature tensor just before the flatten operation is 3 × 3 and we only take the middle anchor point feature, i.e. the one located at(1, 1), whose receptive field does not touch the boundaries of the task image. Based only on that feature vector the network will not be capable to predict the global coordinates of the object observed within that receptive filed. The global coordinates of an object can only be inferred by analysing the information of a subset of all the anchor point features that have a joint receptive field larger than or equal to the whole image. This will also utilise some capacity of the dense layers. But given the x-y meshgrid attached to the task image, even the kernel in the input layer with a receptive field of 3 (assuming a kernel size of 3 × 3) can know where it is, which allows the network to query the global coordinate information of an object locally. In conclusion, the x-y meshgrid is free additional information that can help the task module to capture the relative scale of the scene.

E RENDERING METHOD:

Although the MONet aims to decompose a scene into object-centric representations, the imagined component cannot be rendered independently from the rest parts of the image. This is because the reconstruction outputs of the decoder are the unnormalised logits and directly using softmax to normalise the mask of the imagined component together with the rest parts will constrain the growth of the tool. Imagine that the imagination process that tries to enlarge a tool is assigning higher logits to the area where it thinks the tool should grow towards, while the logits of that area of the background component are still high since the imagination process is not applied to it. During normalisation the imagined component must beat the background to possess the pixels it wants to grow towards and thus the growth is constrained implicitly. We tackle this problem by rather comparing the logits of the imagined tool with the rest components but comparing with itself. We use the minimum logit of the area where the attention mask has a (normalised) probability higher than 0.2 as a threshold before the imagination starts. Any imagined pixel that has a logit higher than that threshold is treated as belonging to the imagined tool. It's common that when people develop a scene decomposition model they normalise the component masks altogether in the end since "sum up to one" is an important inductive bias by design. Nevertheless, our work indicates that this might constrain the down-streaming tasks such as optimising the object-centric representation independently for planning purpose. 

F ADDITIONAL IMAGINATION EXAMPLES:

Figure 6

:

Architecture details are provided in Appendix C.

Figures and Tables:

Figure1: Innovating tools in a simulated reaching task. The model is trained on data triplets comprising the toolkit, the task environment, and an indicator of success (left). The toolkit is an image of stick-like tools (here in yellow, light green, and pink). The goal of the task is to determine if any tool can reach the dark green target whilst touching the red boundary and avoiding any blue obstacles. When no object in the toolkit satisfies these geometric constraints, our model (TasMON) is able to take a tool from the toolkit and imagine suitable modifications until it does satisfy these constraints (right). Here we see TasMON imagine the yellow stick transforming into a hook with a thin enough handle to both pass through the gap between the blue obstacles and reach the dark green target, all whilst still crossing over the red boundary.

Figure2: Illustrative examples of tasks in our dataset. Each task has a green target, a red line indicating the movement limits of the robot, and potentially some blue obstacles. We train on the first four in-sample scenario types: A -D. During evaluation, we also include interpolation scenario types: E -G, which can be solved by the tools seen in training data, as well as an extrapolation scenario type: H, which can only be solved by claw tools never seen in the training data. We choose these scenarios to set constraints on the dimensions of an appropriate tool; for example, in scenario C (third column), a tool needs to be thin enough to reach the goal. Other attributes which affect applicability are tool length and type. In the top row, we overlay tools in yellow, which can press the buttons without colliding with the obstacles or leaving the movement space, as indicated by the green tick. The bottom row, marked with red crosses, shows tasks being attempted by orange tools, which cannot succeed without violating the constraints. Tools are rendered here for visualisation and are not present in the task images.

TasMON model architecture. A convolutional encoder φ represents the task image I G as a latent vector z G . In parallel, the encoder ψ of a MONet scene decomposition model represents the toolkit image I T as a set of latent vectors z T,1..3 where each latent vector represents one of the tools. The concatenated tool-task representations z cat,1..3 are used to compute a soft-attention α 1..3 over tools. A weighted sum z of the tool representations is concatenated to the task representation again forming h cat which is used by a classifier σ to estimate the success of any of the tools solving the task. The task loss L task is computed over σ's output. We use the MONet decoder ψ to render latent tool representations into images. Given the gradient signal produced when optimising task success for the classifier σ, a chosen latent tool representation z im gets updated to render an increasingly suitable tool for the task, as depicted on the bottom right in a callback to Figure 1. A latent vector z 4 for the toolkit background, and a second loss L M ON et used to train the MONet model, are omitted from the figure for clarity of presentation.

Figure 4: Qualitative results from our imagination experiments. The figure shows the evolution of tools during the imagination process (overlayed manually over the corresponding task image). Of these, each row illustrates how the imagination procedure can succeed at constructing tools that solve a task, by (A) increasing the tool length, (B) decreasing the tool width, and (C) changing the configuration between a stick and a hook. The overlapping area is marked in bright yellow. The arrows indicates tool synthesis with increasing optimisation steps.

Figure 6: More examples of imagined tools for all tasks. The first two columns are the task and toolkit inputs, and the next four columns are imagination reconstructions for the selected tool overlaid on the task image. The overlapping area is marked in bright yellow.



Table content: <table><tr><td/><td/><td/><td colspan="2">TI-Infeasible [%]</td><td/><td/><td>TI-Feasible [%]</td><td/></tr><tr><td/><td>n</td><td>RW</td><td colspan="2">FroMON TasMON</td><td>n</td><td>RW</td><td colspan="2">FroMON TasMON</td></tr><tr><td>Scenario</td><td/><td/><td/><td/><td/><td/><td/><td/></tr><tr><td>A</td><td colspan="2">250 65.2%</td><td>87.6%</td><td>92.0%</td><td colspan="2">250 97.2%</td><td>98.8%</td><td>99.2%</td></tr><tr><td>B</td><td colspan="2">250 25.2%</td><td>38.0%</td><td>40.0%</td><td colspan="2">250 87.6%</td><td>90.4%</td><td>84.8%</td></tr><tr><td>C</td><td colspan="2">250 14.4%</td><td>26.4%</td><td>33.2%</td><td colspan="2">250 37.6%</td><td>56.4%</td><td>58.4%</td></tr><tr><td>D</td><td>250</td><td>8.4%</td><td>19.6%</td><td>18.0%</td><td colspan="2">250 28.8%</td><td>30.0%</td><td>37.6%</td></tr><tr><td>E</td><td>250</td><td>8.4%</td><td>20.4%</td><td>16.40%</td><td colspan="2">250 62.8%</td><td>61.2%</td><td>58.8%</td></tr><tr><td>F</td><td>250</td><td>6.0%</td><td>10.8%</td><td>13.2%</td><td colspan="2">250 32.4%</td><td>36.8%</td><td>30.8%</td></tr><tr><td>G</td><td>250</td><td>2.0%</td><td>5.2%</td><td>4.4%</td><td colspan="2">250 25.2%</td><td>27.6%</td><td>24.4%</td></tr><tr><td>H</td><td>250</td><td>0.0%</td><td>0.0%</td><td>0.0%</td><td>250</td><td>0.4%</td><td>0.4%</td><td>0.0%</td></tr><tr><td colspan="3">Overall 2000 16.2%</td><td>26.0%</td><td colspan="3">27.2% 2000 46.5%</td><td>50.2%</td><td>49.3%</td></tr></table>

Please refer to Figure2for illustrative tasks by scenario.

Table content: <table><tr><td/><td colspan="2">TU [%]</td><td>TD [%]</td><td/><td colspan="2">TS [%]</td><td>TI [%]</td><td/></tr><tr><td/><td colspan="8">FroMON TasMON FroMON TasMON FroMON TasMON FroMON TasMON</td></tr><tr><td>Scenario</td><td/><td/><td/><td/><td/><td/><td/><td/></tr><tr><td>A</td><td>95.4%</td><td>97.6%</td><td>97.7%</td><td>97.7%</td><td>95.2%</td><td>96.2%</td><td>83.9%</td><td>87.4%</td></tr><tr><td>B</td><td>88.9%</td><td>91.3%</td><td>97.8%</td><td>97.8%</td><td>82.8%</td><td>79.8%</td><td>48.6%</td><td>48.9%</td></tr><tr><td>C</td><td>83.2%</td><td>85.8%</td><td>97.7%</td><td>97.9%</td><td>88.0%</td><td>88.2%</td><td>38.2%</td><td>40.8%</td></tr><tr><td>D</td><td>77.2%</td><td>79.4%</td><td>97.8%</td><td>97.8%</td><td>78.6%</td><td>75.8%</td><td>20.2%</td><td>20.2%</td></tr><tr><td>E</td><td>76.3%</td><td>77.3%</td><td>97.7%</td><td>97.8%</td><td>72.0%</td><td>71.4%</td><td>29.1%</td><td>24.9%</td></tr><tr><td>F</td><td>37.9%</td><td>36.5%</td><td>97.7%</td><td>97.9%</td><td>62.4%</td><td>65.2%</td><td>15.8%</td><td>16.0%</td></tr><tr><td>G</td><td>48.6%</td><td>48.4%</td><td>97.8%</td><td>97.8%</td><td>76.6%</td><td>73.6%</td><td>14.2%</td><td>12.4%</td></tr><tr><td>H</td><td>37.8%</td><td>38.4%</td><td>96.9%</td><td>97.1%</td><td>69.6%</td><td>70.0%</td><td>0.1%</td><td>0.0%</td></tr><tr><td>Overall</td><td>68.2%</td><td>69.3%</td><td>97.6%</td><td>97.7%</td><td>78.2%</td><td>77.5%</td><td>31.3%</td><td>31.3%</td></tr></table>

presents model performance across all four metrics. Note that the tool-imagination (TI) results here differ from Table1because in this table, we do not warmstart imagination with the same tool for each model, instead letting each model select a tool with highest TU. As a result, models might select different tools to bootstrap imagination, leading to a comparison that is not exactly like-for-like.

Table content: <table/>

Number of instances by scenario type.

Table content: <table><tr><td colspan="3">Type Training Validation</td><td>Test</td></tr><tr><td>A</td><td>10,000</td><td colspan="2">1,000 1,000</td></tr><tr><td>B</td><td>10,000</td><td colspan="2">1,000 1,000</td></tr><tr><td>C</td><td>10,000</td><td colspan="2">1,000 1,000</td></tr><tr><td>D</td><td>10,000</td><td colspan="2">1,000 1,000</td></tr><tr><td>E</td><td>-</td><td colspan="2">-1,000</td></tr><tr><td>F</td><td>-</td><td colspan="2">-1,000</td></tr><tr><td>G</td><td>-</td><td colspan="2">-1,000</td></tr><tr><td>H</td><td>-</td><td colspan="2">-1,000</td></tr><tr><td>Total</td><td>40,000</td><td colspan="2">4,000 8,000</td></tr><tr><td colspan="2">C ARCHITECTURE AND TRAINING DETAILS</td><td/></tr><tr><td>C.1 MODEL ARCHITECTURE</td><td/><td/></tr></table>

Hyperaparmeters used for pretrained MONet.

Table content: <table><tr><td>MONet Hyperparameters</td><td/></tr><tr><td>Attention Blocks</td><td>5</td></tr><tr><td>Attention Channels</td><td>8</td></tr><tr><td>Beta</td><td>0.1</td></tr><tr><td>Gamma</td><td>0.3</td></tr><tr><td colspan="2">Pixel variance background 0.18</td></tr><tr><td>Pixel variance foreground</td><td>0.2</td></tr><tr><td>VAE Channels</td><td>32</td></tr><tr><td>VAE Latents</td><td>16</td></tr><tr><td>D COORDINATE QUERYING</td><td/></tr></table>

Peer review (iclr2020):
The authors constructed an interesting dataset named reaching task, where the model need to predict if the given toolkit is able to solve the corresponding task or not. They showed that combining variational auto-encoding with an auxiliary loss (in this case, the predictor of solving the tasks, could help shaping a latent space where affordance of the tool is encoded as directions in such latent space.) Using the activation maximisation technique (which was phrased as the imagination process), they are able to modify the input tools into the ones that is suitable to solve the corresponding task. I found the idea of using an auxiliary loss when training a VAE may cause the latent space coding direction change novel and interesting. However, I do not find the authors has a strong case of proven it is the case in this manuscript. 1. The performance difference between FroMON and TasMON is not clear. The most critical control model in this paper is the FroMON (frozen MoNet). In this control model, the gradient from the success predictor is not flowing back into the VAE encoder. So, based on the author*s assumption, it should not be benefit of having the tool affordance directions in the latent space. However, in the main results in Table 1. We found the performance between FroMON and TasMON is not quite clear. This is particularly true for the Scenario E, F, G (the interpolation tasks), which is more about generalization and is more important. 2. Are the affordance *directions* in the latent space? The authors used activation maximisation approach to travel in the latent space. My understanding of the approach is it follow the gradient to maximise the predictor*s success prediction in an iterative approach. So, at each optimization step, the z_im can move in different direction. This seems to not fit as a sense of *direction*, as I would assume it is moving along a particular line (not necessarily axis aligned.). Maybe this does explain whey FroMON and TasMon perform equally well. As long as the possible shapes is encoded in a smooth way in the latent space, the activation maximisation could find a path toward the target object. Unfortunately, is that a *direction*? Would it be possible to train an optimization algorithm that is only allow to move in a linear direction, and see how well that work? 3. Why MoNet and multiple tools in the toolkit. A simplified version could drive the point as well. Using MoNet to decompose tools from a toolkit is nice. However, is it really necessary to drive the main point (an auxillary loss of success prediction can shape the latent space of a VAE model) in this paper. In a simplified version, where there is only one tool in the toolkit, one may not need MoNet (maybe still need it for object-background separation?) May the authors comment why multiple tools in the toolkit is important? Minor: 1. typo: page 1, (2nd to the last line). *...that habitual tool use cannot in and off itself ..* --> of 2. A simple video showing how the tool shape change sequentially during the activation maximisation process would be interesting.