# FOCUS: Object-Centric World Models
for Robotics Manipulation

Stefano Ferraro∗ Pietro Mazzaglia∗ Tim Verbelen Bart Dhoedt
Ghent University
Belgium
<EMAIL_ADDRESS>
###### Abstract

Understanding the world in terms of objects and the possible interplays with
them is an important cognition ability, especially in robotics manipulation,
where many tasks require robot-object interactions. However, learning such a
structured world model, which specifically captures entities and
relationships, remains a challenging and underexplored problem. To address
this, we propose FOCUS, a model-based agent that learns an object-centric
world model. Thanks to a novel exploration bonus that stems from the object-
centric representation, FOCUS can be deployed on robotics manipulation tasks
to explore object interactions more easily. Evaluating our approach on
manipulation tasks across different settings, we show that object-centric
world models allow the agent to solve tasks more efficiently and enable
consistent exploration of robot-object interactions. Using a Franka Emika
robot arm, we also showcase how FOCUS could be adopted in real-world settings.

Project website: https://focus-manipulation.github.io/

**footnotetext: These authors contributed equally to this work

##  1 Introduction

For robot manipulators, the tasks we perform as humans are extremely
challenging due to the high level of complexity in the interaction between the
agent and the environment. In recent years, deep reinforcement learning (RL)
has shown to be a promising approach for dealing with these challenging
scenarios [1, 2, 3, 4, 5].

Among RL algorithms, model-based approaches, which allow learning behaviors
using trajectories imagined with the model, a.k.a. ”learning in imagination”
[6], aspire to provide greater data efficiency, compared to the model-free
counterparts [7, 8]. Adopting world models [9, 10], i.e. generative models
that learn the environment dynamics by reconstructing the agent’s
observations, model-based agents have shown impressive performance across
several domains [10, 11, 12], including real-world applications, such as
robotics manipulation and locomotion [13, 14].

However, world models that indistinctly reconstruct all information in the
environment can suffer from several failure modes. For instance, in visual
tasks, they can ignore small, but important features for predicting the
future, such as little objects [15], or they can waste most of the model
capacity on rich, but potentially irrelevant features, such as static
backgrounds [16]. In the case of robot manipulation, this is problematic
because the agent strongly needs to acquire information about the objects to
manipulate in order to solve tasks.

Another challenge in RL for manipulation is engineering reward functions, able
to drive the agent’s learning toward task completion, as attempting to design
dense reward feedback easily leads to faulty reward designs [17, 18, 19, 20].
One solution is to adopt sparse reward feedback, providing a positive reward
only for successful task completion. However, these functions are challenging
to optimize with RL, due to the difficulty of finding rewards in the
environment and thus require appropriate exploration strategies, for which
previous work has resorted to artificial curiosity mechanisms, encouraging the
agent to search for novel interactions with the environment [21, 22].

Humans tend to develop a structured mental model of the world by interacting
with objects, registering specific features associated with objects, such as
shape, color, etc [23]. Since infancy, toddlers learn this by actively
engaging with objects and manipulating them with their hands, discovering
object-centric views that allow them to build an accurate mental model [24,
25].

Inspired by the principle that objects should be of primary importance in the
agent’s world model, we present FOCUS, a model-based RL agent that learns an
object-centric representation of the world and is able to explore object
interactions.

Contributions Our contributions can be summarized as:

  * •

An object-centric latent dynamics model, which learns the dynamics of the
environment, while discriminating object representations into distinct latent
vectors. (Section 3.1);

  * •

An object-centric exploration strategy, which encourages interactions with the
objects, by maximizing the entropy of the latent object representation.
(Section 3.2);

  * •

We show that object-centric world models allow solving dense reward
manipulation tasks faster (Section 4.1);

  * •

We demonstrate that FOCUS’ exploration strategy strongly fosters interaction
with objects and facilitates adaptation to sparse reward tasks, after an
exploration stage (Section 4.2);

  * •

We analyze the effectiveness of the object-centric model, showing that it
captures crucial information for manipulation, both in simulation and in real-
world scenarios (Section 4.3).

##  2 Background

Reinforcement learning. In RL, the agent receives inputs $x$ from the
environment and can interact through actions $a$. The objective of the agent
is to maximize the discounted sum of rewards $\sum_{t}\gamma^{t}r_{t}$, where
$t$ indicates discrete timesteps. To do so, RL agents learn an optimal policy
$\pi(a|x)$ outputting actions that maximize the expected cumulative discounted
reward over time, generally estimated using a critic function, which can be
either a state-value function $v(x)$ or an action-value function $q(x,a)$ [8,
7]. World models [9] additionally learn a generative model of the environment,
capturing the environment dynamics into a latent space, which can be used to
learn the actor and critic functions using imaginary rollouts [10, 12] or to
actively plan at each action [26, 27, 11], which can lead to higher data
efficiency in solving the task.

Exploration Solving sparse-reward tasks is a hard problem in RL. Inspired by
artificial curiosity theories [22, 21], several works have designed
exploration strategies for RL [28, 29, 30]. Other explorations strategies that
have shown great success over time are based upon maximizing uncertainty [31,
32], or the entropy of the agent’s state representation [33, 34, 35]. One
issue with exploration in visual environments is that these approaches can be
particularly attracted by easy-to-reach states that strongly change the visual
appearance of the environment [36]. For robotics manipulation this can cause
undesirable behaviors, e.g. a robot arm exploring different poses in the
proximity of the camera but ignoring interactions with the objects in the
workspace [11].

Object-centric representations. Decomposing scenes into objects can enable
efficient reasoning over high-level building blocks and ensure the agent hones
in on the most relevant concepts [37]. Several 2D object-centric
representations, based on the principle of representing objects separately in
the model, have been studied in the last years [38, 39, 40, 41]. Inspired by
the idea that such representations could help exploit the underlying structure
of our control problem [42, 37], we propose a world model with an object-
centric structured representation [43] that we show could strongly aid
robotics manipulation settings.

##  3 Method

###  3.1 Object-centric World Model

The agent observes the environment through the inputs
$x_{t}=\\{o_{t},q_{t}\\}$ it receives at each interaction, where we can
distinguish the (visual) observations $o_{t}$, e.g. camera RGB and depth, from
the proprioceptive information $q_{t}$, i.e. the robot joint states and
velocities. This information is processed by the agent through an encoder
model, $e_{t}=f_{\phi}(x_{t})$, which can be instantiated as the concatenation
of the outputs of a CNN for high-dimensional observations and an MLP for low-
dimensional proprioceptive inputs.

The world model aims to capture the temporal dynamics of the inputs into a
latent state $s_{t}$. In previous work, this is achieved by reconstructing all
the inputs using an observation decoder. With FOCUS, we are interested in
providing a structured world model representation that separates the object-
specific information into separate latent representations $s_{t}^{obj}$. For
this reason, we instantiate two object-conditioned components: an object
latent extractor and an object decoder.

Figure 1: FOCUS. Visual description of the object-centric world model learned
by FOCUS. (Left) Overall composition of the model, showing the information
flow through the different components. (Right) Detailed view of the decoder,
including the object latent extractor and object decoder.

Models. Overall, the learned world model, which we illustrate in Fig. 1 is
composed of the following components:

$\displaystyle\begin{split}\textrm{Encoder:}&\quad e_{t}=f_{\phi}(x_{t}),\\\
\textrm{Posterior:}&\quad p_{\phi}(s_{t+1}|s_{t},a_{t},e_{t+1}),\\\
\textrm{Prior:}&\quad p_{\phi}(s_{t+1}|s_{t},a_{t}),\\\
\end{split}\begin{split}\textrm{Proprio decoder:}&\quad
p_{\theta}(\hat{q_{t}}|s_{t}),\\\ \textrm{Object latent extractor:}&\quad
p_{\theta}(s_{t}^{obj}|s_{t},c^{obj}),\\\ \textrm{Object decoder:}&\quad
p_{\theta}(\hat{o}_{t}^{obj},w_{t}^{obj}|s_{t}^{obj}),\end{split}$ (1)

where $\phi$ indicates the set of learnable parameters.

We adopt a recurrent state-space model (RSSM) [44] for the dynamics
components, i.e. prior and posterior, which extracts a latent state $s_{t}$
made of a deterministic and a stochastic component. Proprioceptive information
$\hat{q_{t}}$ is decoded out of the latent state $s_{t}$, using an MLP. For
each object in the scene, the object latent extractor receives the world
model’s latent state $s_{t}$ and a vector identifying the object $c^{obj}$,
and extracts an object-centric latent $s_{t}^{obj}$. Given such object latent
vectors, the object decoder outputs two values: the object-related observation
information $\hat{o}_{t}^{obj}$, and the segmentation “unnormalized weights”
$w^{obj}_{t}$.

The object-related information $\hat{o}_{t}^{obj}$ is a masked version of the
visual inputs $o_{t}$ where all the irrelevant information to the object is
masked out, using a segmentation mask. The one-dimensional “unnormalized
weights” $w^{obj}_{t}$ represent object-specific per-pixel logits for
segmenting the objects. The overall scene segmentation is obtained by applying
a softmax among all object weights, where each pixel is assigned to the object
that outputs the highest weight.

In practice, the object latent extractor is instantiated as an MLP while the
object decoder is instantiated as a transposed CNN. For the object-
conditioning vector $c^{obj}$ we adopt a one-hot vector identifying the object
instance.

Objective. The model is trained end-to-end by minimizing the following loss:

$\mathcal{L_{\textrm{wm}}}=\mathcal{L}_{\textrm{dyn}}+\mathcal{L}_{\textrm{proprio}}+\mathcal{L}_{\textrm{obj}}.$
(2)

The dynamics minimizes the Kullback–Leibler (KL) divergence between posterior
and prior:

$\mathcal{L_{\textrm{dyn}}}=D_{\textrm{KL}}[p_{\phi}(s_{t+1}|s_{t},a_{t},e_{t+1})||p_{\phi}(s_{t+1}|s_{t},a_{t})].$
(3)

The proprioceptive decoder learns to reconstruct proprio states, by minimizing
a negative log-likelihood (NLL) loss:

$\begin{split}\mathcal{L}_{\textrm{proprio}}=-\log
p_{\theta}(\hat{q_{t}}|s_{t})\end{split}$ (4)

The object decoder learns to reconstruct object-centric information,
outputting “object weights” for the segmentation mask and reconstructing
object-related observations. The observation is masked via an object-specific
mask $m^{i}_{t}$, to focus only on the $i$-th object information in the loss.
The object decoder loss can be expressed as:

$\begin{split}\mathcal{L_{\textrm{obj}}}=-\log\underbrace{p(\hat{m_{t}})}_{\textrm{mask}}-\log\sum^{N}_{i=0}\underbrace{m_{t}^{i}p_{\theta}(\hat{x_{t}^{i}}|s_{t}^{i})}_{\textrm{masked
reconstruction}}\end{split}$ (5)

where the overall segmentation mask is obtained as:

$\hat{m_{t}}=\textrm{softmax}(w^{1}_{t},...,w^{N}_{t})$ (6)

with $N$ being the number of all object instances present in the scene111The
scene, with objects masked out, is also considered a ”special object”.. By
minimizing the masked reconstruction NLL, the object-decoder ensures that each
object latent $s^{i}$ focuses on capturing only its relevant information, as
the reconstructions obtained from the latent are masked per object.
Furthermore, objects compete for occupying their correct space in the scene
(in pixel space), through the mask loss. Thus, each object’s latent vector is
aware of the object’s positioning in the scene.

In addition to the above components, when the agent intends to solve reward-
supervised tasks, a reward predictor $p_{\phi}(\hat{r}_{t}|s_{t})$ is learned
to minimize the loss $\mathcal{L_{\textrm{rew}}}=-\log
p_{\phi}(\hat{r}_{t}|s_{t})$. Further details about the model architecture,
the objective, and the optimization process are provided in Appendix.

###  3.2 Object-centric Exploration

State maximum entropy approaches for RL [35, 34, 33] learn an environment
representation, on top of which they compute an entropy estimate that is
maximized by the agent’s actor to foster exploration of diverse environment
states. Given our object-centric representation, we can incentivize
exploration towards object interactions and discovery of novel object views,
by having the agent maximize the entropy over the object latent vectors.

In order to estimate the entropy value over training batches, we apply a K-NN
particle-based estimator [45] on top of the object latent representation. By
maximizing the overall entropy, with respect to all objects in the scene, we
derive the following reward for object-centric exploration:

$r_{\textrm{expl}}=\sum^{N}_{obj=0}r_{\textrm{expl}}^{obj},\quad\textrm{where}\quad
r_{\textrm{expl}}^{obj}(s)\propto\sum^{K}_{i=1}\log\mathinner{\\!\left\lVert
s^{obj}-s^{obj}_{i}\right\rVert}_{2}$ (7)

where $s^{obj}$ is extracted from $s$ using the object latent extractor,
conditioned by the vector $c^{obj}$, and $s^{obj}_{i}$ is the $i$-th nearest
neighbor to $s^{obj}$.

Crucially, as we use the world model to optimize actor and critic networks in
imagination [10], the latent states in Equation 7 are states of imaginary
trajectories, generated by the world model by following the actor’s proposed
actions.

Figure 2: Manipulation tasks. Comparing FOCUS and Dreamer over dense
supervised manipulation tasks, to assess the quality of the world model. (3
seeds)

##  4 Experiments

We empirically evaluate FOCUS aiming to answer the following questions:

  * •

Does FOCUS object-centric world model enable faster learning of robotics
manipulation tasks? (Section 4.1) We test FOCUS against the DreamerV2 world
model-based agent on a set of supervised dense reward tasks.

  * •

Can our object-driven exploration lead to more meaningful exploration, i.e.
exploring interactions with the objects in the scene? (Section 4.2) We deploy
different exploration strategies across several environments and compare them
using exploration metrics. Then, we evaluate whether the exploration stage was
useful, by fine-tuning the exploration approaches on sparse-reward tasks in
the same environments.

  * •

Does the learned world model captures crucial information for manipulation
better? (Section 4.3) We qualitatively assess the representation learned by
FOCUS, both in simulation and real-world environments.

Simulation environments. We tested our approach in two simulation
environments: ManiSkill2 [46] and robosuite [47]. For the robosuite
environment, we considered a single (red) cube and a two-cube (red and green)
setup. For Maniskill2, we opted for two single YCB objects setups (banana and
master chef can), and a faucet setup. Segmentation masks for training FOCUS’
decoder are provided by the simulator. Find the environments illustrated in
Appendix.

###  4.1 Supervised dense reward tasks

In Figure 2, we compare the performance in terms of task episodic rewards that
are obtained by FOCUS compared to DreamerV2 on supervised dense-reward
manipulation tasks.

As both agents are world model-based agents that learn actor and critic in
imagination (and we do not use the object-centric exploration for FOCUS in
these experiments), the only difference between the two methods is how the
world model is learned, so this study compares the quality of the world model
to learn control policies.

We observe that FOCUS converges faster in the Lift Cube environment and
converges to higher rewards in the Stack Cube ($\sim$+10%) and Turn Faucet
($\sim$+35%) tasks. This provides empirical evidence that the object-centric
model facilitates the learning of dense-reward task behaviors.

###  4.2 Exploration and sparse reward tasks

For each environment, we consider a sparse-reward task setting. The tasks
analyzed are either “lift” (L) or “push” (P) tasks, for cube and YCB objects
environments, and “turn” (T) for the faucet. The two-stage experimental setup
is inspired by the unsupervised RL benchmark [48].

During the exploration stage, all exploration approaches (excl. Dreamer), can
find the rewards but they do not maximize the environment rewards yet, as they
only focus on exploration, using their own exploration rewards. The
exploration stage lasts for 2M environment steps.

During the adaptation stage, the agents, pre-trained on the data they
collected during the exploration stage, are allowed additional environment
steps, aiming to solve the sparse reward tasks, maximizing the reward
function. This stage evaluates whether the exploration data was useful to
solve downstream tasks.

Exploration stage. To evaluate the performance of exploration approaches for
robotics manipulation, we look into exploration metrics that are related to
the interaction with the object and the ability to find potentially rewarding
areas in the environment. The metrics we adopt are:

  * •

Contact (%): average percentage of contact interactions between the gripper
and the objects over an episode.

  * •

Positional and angular displacement: cumulative position (m) and angular (rad)
displacement of all the objects over an entire episode.

  * •

Up, far, close, left, right placement: number of times the object is moved in
the relative area of the working space. Details about how the areas are
defined are in Appendix.

| Contact | Pos Disp | Ang Disp | Placement
---|---|---|---|---
Up | ↑ | ↓ | ← | →
Red cube | FOCUS | 0.04 | 0.15 | 1.86 | 3.5 | 37.45 | 0.06 | 0.15 | 0.01
Dreamer | 0.0 | 0.0 | 0.0 | 0.0 | 0.0 | 0.0 | 0.0 | 0.0
P2E | 0.0 | 0.0 | 0.01 | 0.0 | 0.15 | 0.0 | 0.0 | 0.0
APT | 0.0 | 0.01 | 0.15 | 0.01 | 0.11 | 0.0 | 0.0 | 0.0
RG cubes | FOCUS | 0.1 | 0.4 | 4.46 | 0.0 | 0.0 | 0.53 | 0.0 | 0.0
Dreamer | 0.0 | 0.01 | 0.01 | 0.0 | 0.0 | 0.16 | 0.0 | 0.0
P2E | 0.0 | 0.01 | 0.01 | 0.0 | 0.0 | 0.3 | 0.0 | 0.0
APT | 0.01 | 0.03 | 0.25 | 0.06 | 0.0 | 5.05 | 0.0 | 0.03
Faucet* | FOCUS | 0.12 | - | 1.22 | - | - | - | - | -
Dreamer | 0.0 | - | 0.01 | - | - | - | - | -
P2E | 0.1 | - | 1.23 | - | - | - | - | -
APT | 0.08 | - | 0.98 | - | - | - | - | -
Banana | FOCUS | 0.08 | 0.38 | 3.5 | 1.31 | 25.28 | 0.72 | 6.08 | 7.12
Dreamer | 0.0 | 0.0 | 0.06 | 0.04 | 0.22 | 0.02 | 0.16 | 0.29
P2E | 0.02 | 0.11 | 0.97 | 0.64 | 8.48 | 0.22 | 4.78 | 4.14
APT | 0.01 | 0.04 | 0.4 | 0.24 | 2.0 | 0.98 | 1.3 | 2.1
Master Chef Can | FOCUS | 0.05 | 0.49 | 5.78 | 0.17 | 145.66 | 28.57 | 61.51 | 62.8
Dreamer | 0.02 | 0.26 | 3.04 | 0.08 | 85.68 | 24.02 | 24.34 | 46.35
P2E | 0.03 | 0.29 | 3.26 | 0.12 | 69.95 | 30.63 | 29.72 | 36.0
APT | 0.01 | 0.14 | 1.56 | 0.06 | 23.64 | 23.74 | 15.88 | 13.52

Table 1: Exploration stage results. Exploration metrics across different
environments. *Position displacement and placements are absent for the Faucet
as the object is not movable.

In Table 1, we compare FOCUS against two exploration strategies: Plan2Explore
(P2E) [32] and a combination of the Dreamer algorithm and Active Pre-training
(APT) [33, 11]. We also test a DreamerV2 agent, which aims to maximize sparse
rewards in the environment 222DreamerV2 is still able to explore, with limited
capacity, thanks to the actor entropy term in the actor’s objective [10]..

We find that FOCUS shines across all environments, showing strong object
interaction performance, and outperforming all other approaches in most
metrics. Dreamer manages to find sparse rewards, and thus explore object
interactions consistently, only in the master chef can ‘push’ environment, but
struggles elsewhere. APT and P2E tend to explore object interactions more
consistently than Dreamer, but generally move objects less than FOCUS and tend
to find potentially rewarding areas, i.e. achieving placements, more rarely.
Detailed results, showing exploration metrics over timesteps, are presented in
Appendix and visual results of FOCUS exploration policies are provided in the
project website.

Adaptation stage. After exploring the environment for 2M environment steps, we
adapt the exploration approaches, allowing them an additional number of
environment interactions. During this stage, if the exploration stage was
fruitful, the agent may already know where the rewards can be found in the
environment, despite not having actively attempted to maximize them.

In order to speed up the adaptation stage, task-driven actor and critic
networks are pre-trained in imagination, during the exploration stage, using
the reward predictor trained on the exploration stage rewards. These actor and
critic networks are then fine-tuned during the adaptation stage. This setup is
similar to the few-shot adaptation experiments presented in [32].

To provide an idea of how frequently the agents were able to find sparse
rewards during the exploration stage, we provide aggregate metrics in Table 2.
The adaptation curves, showing episode rewards over time, are presented in
Figure 3. Results clearly show that FOCUS is the only method that makes
significant progress across all tasks. As the figures in the Table show, this
is mainly due to the fact that FOCUS was also the approach that more
frequently found rewards during the exploration stage, making adaptation to
downstream tasks easier.

| Red Cube Push | Red Green Cube Lift | Faucet Turn | Banana Push | Banana Lift | MC Can Push | MC Can Lift
---|---|---|---|---|---|---|---
FOCUS | 0.67 | 0.0 | 101.67 | 38.67 | 27.0 | 507.67 | 2.67
Dreamer | 0.0 | 0.0 | 0.0 | 2.67 | 0.67 | 628.67 | 0.33
P2E | 0.0 | 0.0 | 35.33 | 19.67 | 14.67 | 207.67 | 2.33
APT | 0.0 | 2.0 | 34.33 | 6.67 | 5.33 | 68.33 | 1.67

Table 2: Sparse rewards through exploration. Summarizing how many times the
agents were able to find the sparse rewards during the exploration stage. (3
seeds) Figure 3: Adaptation stage. Comparing adaptation performance of FOCUS
and other exploration strategies in sparse-reward manipulation tasks. (3
seeds)

((a)) Red cube

((b)) RG cubes

((c)) Banana

((d)) Master chef can

Figure 4: Simulation reconstructions. Comparison of Dreamer and FOCUS
reconstructions on random action trajectories.

###  4.3 Model Analysis

Simulation. We present results in the simulated environments, sampling random
action trajectories, in Figure 4. In order to qualitatively assess the
capacity of FOCUS to better capture information about the objects, we compare
reconstructions from FOCUS to the non-object-centric world model of DreamerV2.
We observe that, while Dreamer can reconstruct the robot and the background,
it struggles to accurately reconstruct objects. The banana is reduced to a
yellow spot, while the can is almost invisible against the background. FOCUS,
instead, reconstructs the objects more accurately, by capturing the object
information in the object latents.

Real-world. In order to show that FOCUS can also be applied to real-world
setups, we present results in a real-world setting, where we use a Franke
Emika arm in front of a white table with three colored bricks (blue, orange,
red). The main issue with applying FOCUS in the real world comes from the
absence of segmentation masks. However, thanks to recent progress in the study
of large-scale vision models, the issue can be easily circumvented using a
pre-trained segmentation model. For these experiments, we adopted the Segment
Anything Model (SAM) [49], which we found particularly proficient at providing
segmentation masks and tracking objects, after providing a single labeled
sample to recognize the bricks in the scene. In Figure 5, we show that our
method effectively captures object information, though sometimes sacrificing
information about the gripper. We saw this happening for the simulated
environments as well (Figure 4), nonetheless, this did not affect performance
negatively. We believe the reason for that is that the agent can exploit the
proprioceptive information for control. However, if the agent does not capture
information about objects correctly, this can significantly harm the task’s
performance.

Figure 5: Real world reconstructions. Comparing ground truth and FOCUS
reconstructions on real-world sequences of data.

##  5 Conclusion

We presented FOCUS, an object-centric model-based agent for robotics
manipulation. Our results show that our approach can learn manipulation tasks
more efficiently, thanks to the object-centric world model, and explore robot-
object interactions more easily, thanks to a novel object-driven exploration
reward. As we showed, this can facilitate solving even challenging sparse
rewards tasks, where finding rewards can be hard and require demonstrations or
task-specific strategies. In the future, we aim to develop our object-centric
representation further, aiming to remove dependence upon the segmentation
information, used in the object decoder’s loss, which is not easily available
to a real robot. In order to overcome such limitations, we aim to investigate
unsupervised strategies for scene decomposition [50, 38].

##  6 Limitations

FOCUS requires segmentation masks over the object of attention in order to
function. Segmentation masks can be challenging to obtain in real-world
scenarios. To overcome this limitation, we developed a segmentation pipeline
based on SAM [51, 52].

In our current implementation, every object of interest is uniquely identified
by a one-hot vector $c^{obj}$. This requires prior knowledge about the number
and the identity of objects in the scene. In the future, this should be
addressed by relying on more sophisticated sets of object features, which
would allow us to identify objects, before feeding the information to the
object decoder.

Reconstruction-based models may struggle with changing dynamics in the
environment or complex backgrounds [53], making the adoption of world models
challenging in real-world scenarios, where the environment is often
uncontrolled. To address this limitation in the future, we aim to remove the
background from the reconstruction scope, focusing solely on the objects in
the scene and on the proprioceptive information, or using some alternative
learning paradigms, e.g. contrastive learning [54], to reduce the amount of
information stored in the representation about the background.

#### Acknowledgments

This research received funding from the Flemish Government (AI Research
Program). Pietro Mazzaglia is funded by a Ph.D. grant of the Flanders Research
Foundation (FWO).

## References

  * Levine et al. [2016] S. Levine, C. Finn, T. Darrell, and P. Abbeel.  End-to-end training of deep visuomotor policies.  _J. Mach. Learn. Res._ , 2016.
  * OpenAI et al. [2019] OpenAI, I. Akkaya, M. Andrychowicz, M. Chociej, M. Litwin, B. McGrew, A. Petron, A. Paino, M. Plappert, G. Powell, R. Ribas, J. Schneider, N. A. Tezak, J. Tworek, P. Welinder, L. Weng, Q. Yuan, W. Zaremba, and L. M. Zhang.  Solving rubik’s cube with a robot hand.  _ArXiv_ , abs/1910.07113, 2019.
  * Kalashnikov et al. [2018] D. Kalashnikov, A. Irpan, P. Pastor, J. Ibarz, A. Herzog, E. Jang, D. Quillen, E. Holly, M. Kalakrishnan, V. Vanhoucke, and S. Levine.  Qt-opt: Scalable deep reinforcement learning for vision-based robotic manipulation.  _ArXiv_ , abs/1806.10293, 2018.
  * Lu et al. [2021] Y. Lu, K. Hausman, Y. Chebotar, M. Yan, E. Jang, A. Herzog, T. Xiao, A. Irpan, M. Khansari, D. Kalashnikov, and S. Levine.  AW-opt: Learning robotic skills with imitation andreinforcement at scale.  In _5th Annual Conference on Robot Learning (CoRL)_ , 2021.
  * Lee et al. [2021] A. X. Lee, C. Devin, Y. Zhou, T. Lampe, K. Bousmalis, J. T. Springenberg, A. Byravan, A. Abdolmaleki, N. Gileadi, D. Khosid, C. Fantacci, J. E. Chen, A. S. Raju, R. Jeong, M. Neunert, A. Laurens, S. Saliceti, F. Casarini, M. A. Riedmiller, R. Hadsell, and F. Nori.  Beyond pick-and-place: Tackling robotic stacking of diverse shapes.  _ArXiv_ , abs/2110.06192, 2021.
  * Hafner et al. [2020] D. Hafner, T. Lillicrap, J. Ba, and M. Norouzi.  Dream to control: Learning behaviors by latent imagination, 2020.
  * Fujimoto et al. [2018] S. Fujimoto, H. van Hoof, and D. Meger.  Addressing function approximation error in actor-critic methods, 2018\.
  * Haarnoja et al. [2018] T. Haarnoja, A. Zhou, P. Abbeel, and S. Levine.  Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with a stochastic actor, 2018.
  * Ha and Schmidhuber [2018] D. Ha and J. Schmidhuber.  World models.  2018\.  doi:10.5281/ZENODO.1207631.  URL https://zenodo.org/record/1207631.
  * Hafner et al. [2021] D. Hafner, T. P. Lillicrap, M. Norouzi, and J. Ba.  Mastering atari with discrete world models.  In _ICLR_ , 2021.
  * Rajeswar et al. [2023] S. Rajeswar, P. Mazzaglia, T. Verbelen, A. Piché, B. Dhoedt, A. Courville, and A. Lacoste.  Mastering the unsupervised reinforcement learning benchmark from pixels.  2023\.
  * Hafner et al. [2023] D. Hafner, J. Pasukonis, J. Ba, and T. Lillicrap.  Mastering diverse domains through world models.  _arXiv preprint arXiv:2301.04104_ , 2023.
  * Wu et al. [2022] P. Wu, A. Escontrela, D. Hafner, K. Goldberg, and P. Abbeel.  Daydreamer: World models for physical robot learning, 2022.
  * Seo et al. [2023] Y. Seo, J. Kim, S. James, K. Lee, J. Shin, and P. Abbeel.  Multi-view masked world models for visual robotic manipulation, 2023.
  * Seo et al. [2022] Y. Seo, D. Hafner, H. Liu, F. Liu, S. James, K. Lee, and P. Abbeel.  Masked world models for visual control, 2022.
  * Deng et al. [2021] F. Deng, I. Jang, and S. Ahn.  Dreamerpro: Reconstruction-free model-based reinforcement learning with prototypical representations, 2021.
  * Amodei et al. [2016] D. Amodei, C. Olah, J. Steinhardt, P. Christiano, J. Schulman, and D. Mané.  Concrete problems in ai safety, 2016.
  * Clark and Amodei [2016] J. Clark and D. Amodei.  Faulty reward functions in the wild.  https://openai.com/blog/faulty-reward-functions/, 2016.  Accessed: 2022-04-19.
  * Krakovna et al. [2020] V. Krakovna et al.  Specification gaming: the flip side of ai ingenuity.  https://www.deepmind.com/blog/specification-gaming-the-flip-side-of-ai-ingenuity, 2020\.  Accessed: 2022-04-19.
  * Popov et al. [2017] I. Popov et al.  Data-efficient deep reinforcement learning for dexterous manipulation, 2017.
  * Oudeyer et al. [2007] P. Oudeyer, F. Kaplan, and V. V. Hafner.  Intrinsic motivation systems for autonomous mental development.  _IEEE Transactions on Evolutionary Computation_ , 11(2):265–286, 2007.
  * Schmidhuber [1991] J. Schmidhuber.  Curious model-building control systems.  In _[Proceedings] 1991 IEEE International Joint Conference on Neural Networks_ , pages 1458–1463 vol.2, 1991.
  * Hawkins et al. [2017] J. Hawkins, S. Ahmad, and Y. Cui.  A theory of how columns in the neocortex enable learning the structure of the world.  _Frontiers in Neural Circuits_ , 11, 2017.  ISSN 1662-5110.  doi:10.3389/fncir.2017.00081.  URL https://www.frontiersin.org/articles/10.3389/fncir.2017.00081.
  * Smith et al. [2018] L. B. Smith, S. Jayaraman, E. Clerkin, and C. Yu.  The developing infant creates a curriculum for statistical learning.  _Trends in Cognitive Sciences_ , 22(4):325–336, 2018.  ISSN 1364-6613.  doi:https://doi.org/10.1016/j.tics.2018.02.004.  URL https://www.sciencedirect.com/science/article/pii/S1364661318300275.
  * Slone et al. [2019] L. Slone, L. Smith, and C. Yu.  Self‐generated variability in object images predicts vocabulary growth.  _Developmental Science_ , 22, 02 2019.  doi:10.1111/desc.12816.
  * Schrittwieser et al. [2020] J. Schrittwieser, I. Antonoglou, T. Hubert, K. Simonyan, L. Sifre, S. Schmitt, A. Guez, E. Lockhart, D. Hassabis, T. Graepel, T. Lillicrap, and D. Silver.  Mastering atari, go, chess and shogi by planning with a learned model.  _Nature_ , 588(7839):604–609, dec 2020.  doi:10.1038/s41586-020-03051-4.  URL https://doi.org/10.1038%2Fs41586-020-03051-4.
  * Hansen et al. [2022] N. Hansen, X. Wang, and H. Su.  Temporal difference learning for model predictive control, 2022.
  * Pathak et al. [2017] D. Pathak, P. Agrawal, A. A. Efros, and T. Darrell.  Curiosity-driven exploration by self-supervised prediction, 2017.
  * Mazzaglia et al. [2022] P. Mazzaglia, O. Catal, T. Verbelen, and B. Dhoedt.  Curiosity-driven exploration via latent bayesian surprise, 2022.
  * Rajeswar et al. [2021] S. Rajeswar, C. Ibrahim, N. Surya, F. Golemo, D. Vazquez, A. Courville, and P. O. Pinheiro.  Touch-based curiosity for sparse-reward tasks, 2021.
  * Pathak et al. [2019] D. Pathak, D. Gandhi, and A. Gupta.  Self-supervised exploration via disagreement, 2019.
  * Sekar et al. [2020] R. Sekar, O. Rybkin, K. Daniilidis, P. Abbeel, D. Hafner, and D. Pathak.  Planning to explore via self-supervised world models.  In _ICML_ , 2020.
  * Liu and Abbeel [2021] H. Liu and P. Abbeel.  Behavior from the void: Unsupervised active pre-training.  In M. Ranzato, A. Beygelzimer, Y. Dauphin, P. Liang, and J. W. Vaughan, editors, _Advances in Neural Information Processing Systems_ , volume 34, pages 18459–18473. Curran Associates, Inc., 2021.
  * Seo et al. [2021] Y. Seo, L. Chen, J. Shin, H. Lee, P. Abbeel, and K. Lee.  State entropy maximization with random encoders for efficient exploration, 2021.
  * Mutti et al. [2021] M. Mutti, L. Pratissoli, and M. Restelli.  Task-agnostic exploration via policy gradient of a non-parametric state entropy estimate, 2021.
  * Burda et al. [2018] Y. Burda, H. Edwards, D. Pathak, A. Storkey, T. Darrell, and A. A. Efros.  Large-scale study of curiosity-driven learning, 2018.
  * Dittadi et al. [2022] A. Dittadi, S. Papa, M. D. Vita, B. Schölkopf, O. Winther, and F. Locatello.  Generalization and robustness implications in object-centric learning, 2022.
  * Locatello et al. [2020] F. Locatello, D. Weissenborn, T. Unterthiner, A. Mahendran, G. Heigold, J. Uszkoreit, A. Dosovitskiy, and T. Kipf.  Object-centric learning with slot attention, 2020.
  * Greff et al. [2020] K. Greff, R. L. Kaufman, R. Kabra, N. Watters, C. Burgess, D. Zoran, L. Matthey, M. Botvinick, and A. Lerchner.  Multi-object representation learning with iterative variational inference, 2020.
  * Burgess et al. [2019] C. P. Burgess, L. Matthey, N. Watters, R. Kabra, I. Higgins, M. Botvinick, and A. Lerchner.  Monet: Unsupervised scene decomposition and representation, 2019.
  * Nakano et al. [2023] A. Nakano, M. Suzuki, and Y. Matsuo.  Interaction-based disentanglement of entities for object-centric world models.  In _The Eleventh International Conference on Learning Representations_ , 2023.  URL https://openreview.net/forum?id=JQc2VowqCzz.
  * Diuk et al. [2008] C. Diuk, A. Cohen, and M. L. Littman.  An object-oriented representation for efficient reinforcement learning.  In _Proceedings of the 25th international conference on Machine learning_ , pages 240–247, 2008.
  * Kipf et al. [2020] T. Kipf, E. van der Pol, and M. Welling.  Contrastive learning of structured world models, 2020.
  * Hafner et al. [2019] D. Hafner, T. Lillicrap, I. Fischer, R. Villegas, D. Ha, H. Lee, and J. Davidson.  Learning latent dynamics for planning from pixels.  In _ICML_ , pages 2555–2565, 2019.
  * Singh et al. [2003] H. Singh, N. Misra, V. Hnizdo, A. Fedorowicz, and E. Demchuk.  Nearest neighbor estimates of entropy.  _American journal of mathematical and management sciences_ , 23(3-4):301–321, 2003.
  * Gu et al. [2023] J. Gu, F. Xiang, X. Li, Z. Ling, X. Liu, T. Mu, Y. Tang, S. Tao, X. Wei, Y. Yao, X. Yuan, P. Xie, Z. Huang, R. Chen, and H. Su.  Maniskill2: A unified benchmark for generalizable manipulation skills, 2023.
  * Zhu et al. [2020] Y. Zhu, J. Wong, A. Mandlekar, R. Martín-Martín, A. Joshi, S. Nasiriany, and Y. Zhu.  robosuite: A modular simulation framework and benchmark for robot learning.  In _arXiv preprint arXiv:2009.12293_ , 2020.
  * Laskin et al. [2021] M. Laskin, D. Yarats, H. Liu, K. Lee, A. Zhan, K. Lu, C. Cang, L. Pinto, and P. Abbeel.  Urlb: Unsupervised reinforcement learning benchmark, 2021.
  * Kirillov et al. [2023] A. Kirillov, E. Mintun, N. Ravi, H. Mao, C. Rolland, L. Gustafson, T. Xiao, S. Whitehead, A. C. Berg, W.-Y. Lo, P. Dollár, and R. Girshick.  Segment anything, 2023.
  * Wu et al. [2023] Z. Wu, N. Dvornik, K. Greff, T. Kipf, and A. Garg.  Slotformer: Unsupervised visual dynamics simulation with object-centric models, 2023.
  * Kirillov et al. [2023] A. Kirillov, E. Mintun, N. Ravi, H. Mao, C. Rolland, L. Gustafson, T. Xiao, S. Whitehead, A. C. Berg, W.-Y. Lo, et al.  Segment anything.  _arXiv preprint arXiv:2304.02643_ , 2023.
  * Cheng et al. [2023] Y. Cheng, L. Li, Y. Xu, X. Li, Z. Yang, W. Wang, and Y. Yang.  Segment and track anything.  _arXiv preprint arXiv:2305.06558_ , 2023.
  * Stone et al. [2021] A. Stone, O. Ramirez, K. Konolige, and R. Jonschkowski.  The distracting control suite – a challenging benchmark for reinforcement learning from pixels, 2021.
  * Chen et al. [2020] T. Chen, S. Kornblith, M. Norouzi, and G. Hinton.  A simple framework for contrastive learning of visual representations, 2020.
  * Cho et al. [2014] K. Cho, B. van Merrienboer, D. Bahdanau, and Y. Bengio.  On the properties of neural machine translation: Encoder-decoder approaches, 2014.  URL https://arxiv.org/abs/1409.1259.

##  Appendix A Model details

Dreamer. The architecture adopted for Dreamer is relevant to the one
documented by [10]. Model states have both a deterministic and a stochastic
component: the deterministic component is the 200-dimensional output of a GRU
([55], with a 200-dimensional hidden layer; the stochastic component consists
of 32 categorical distributions with 32 classes each. States-based inputs such
as the proprioception, the encoder, and the decoder are 4-layer MLP with a
dimensionality of 400. For pixels-based inputs, the encoder and decoder follow
the architecture of Dreamer [10], taking 64 × 64 RGBD images as inputs. Both
encoder and decoder networks have a depth factor of 48. To ensure stable
training during the initial phases, we adopt a technique from [12] where the
weights of the output layer in the critic network are initialized to zero.
This approach contributes to the stabilization of the training process
especially in the early stages of training.

Networks are updated by sampling batches of 32 sequences of 32 timesteps,
using Adam with learning rate $3\mathrm{e}^{-4}$ for the updates, and clipping
gradients norm to 100.

FOCUS. The architecture proposed is based on the implementation of Dreamer
described above. The encoding network and the state-based decoding unit have
the same structure mentioned in Dreamer. We introduced an object latent
extractor unit consisting of a 3-layer MLP with a dimensionality of 512. The
object-decoder network resembles the structure of the Dreamer’s decoder, the
depth factor for the CNN is set to 72. 64x64 RGBD images along with a
”segmentation weights” image are generated per each object.

Objective. In the objective of FOCUS we describe object-specific masks
$m^{i}_{t}$. These masks are binary images, obtained from the entire scene
segmentation mask. In practice,
$m_{t}=\sum_{i}{m^{i}_{t}}\cdot\arg\max{c^{i}}$, which means the overall
segmentation mask is obtained by summing the object-specific masks, multiplied
by their object index, which can be extracted from $c^{i}$ through an
$\arg\max$, being $c^{i}$ a one-hot vector.

##  Appendix B Environment details

((a)) Red cube

((b)) RG cubes

((c)) Faucet

((d)) Banana

((e)) Master Chef Can

Figure 6: Environments. Visualization of the environments we applied FOCUS on.
(a-b) from robosuite. (c-e) from ManiSkill2.

In both simulation settings, ManiSkill and robosuite, the robotic agent is a
Franka arm with 7-DoF, controlled in cartesian end-effector space with fixed
gripper orientation. Gripper state (open/close) adds a degree of freedom. For
all the mentioned environments objects are spawned at a fixed pose. Visual
observations are rendered at 64x64 resolution for all the environments.
Considered tasks are shown in figure 6.

###  B.1 Metrics and task rewards

To evaluate the placement metric, we divide the workspace in front of the
robot into 5 areas, with respect to the 5 cited directions: right, left, far,
close, and up. For the Robosuite environment, with respect to the origin of
the environment (coinciding with the center of the table) we consider a
successful placement if the object is placed at a minimum of 0.25m up to a
maximum of 0.4m along the directions considered. For the ”up” placement the
minimum threshold is reduced to 0.05m.

For the Maniskill2 environment, in the same logic, we considered a minimum
threshold of 0.4m up to a maximum of 0.5m. For the ”up” placement a minimum
threshold is reduced to 0.1m.

The definition of the task rewards for the exploration tasks is linked to the
placement definition given. Any reward is provided till the object has been
placed in the desired area. For the lift tasks, before assigning any reward we
check if the object is grasped. The definition of the task rewards for dense
tasks is the one defined by the default implementation.

##  Appendix C Extended results

((a)) Red cube

((b)) RG cubes

((c)) Faucet

((a)) Banana

((b)) Master Chef Can

Figure 8: Model reconstructions over random actions for simulated
environments. First row is ground truth. Second row is FOCUS reconstruction,
overlap of instance masks. Last row shows Dreamer reconstruction.

Figure 9: Displaying exploration metrics over time for the exploration
experiments.