Original: Curious Exploration via Structured World Models Yields Zero-Shot Object Manipulation
Revised: Curious Exploration via Structured World Models Yields Zero-Shot Object Manipulation

Original: Abstract
Revised: Abstract

Original: It has been a long-standing dream to design artificial agents that explore theirenvironment efficiently via intrinsic motivation, similar to how children performcurious free play.
Revised: It has been a long-standing dream to design artificial agents that explore their environment efficiently via intrinsic motivation, similar to how children perform curious free play.

Original: Despite recent advances in intrinsically motivated reinforcement learning (RL), sample-efficient exploration in object manipulation scenariosremains a significant challenge as most of the relevant information lies in thesparse agent-object and object-object interactions.
Revised: Despite recent advances in intrinsically motivated reinforcement learning (RL), sample-efficient exploration in object manipulation scenarios remains a significant challenge as most of the relevant information lies in the sparse agent-object and object-object interactions.

Original: In this paper, we propose to usestructured world models to incorporate relational inductive biases in the controlloop to achieve sample-efficient and interaction-rich exploration in compositionalmulti-object environments.
Revised: In this paper, we propose to use structured world models to incorporate relational inductive biases in the control loop to achieve sample-efficient and interaction-rich exploration in compositional multi-object environments.

Original: By planning for future novelty inside structured worldmodels, our method generates free-play behavior that starts to interact with objectsearly on and develops more complex behavior over time.
Revised: By planning for future novelty inside structured world models, our method generates free-play behavior that starts to interact with objects early on and develops more complex behavior over time.

Original: Instead of using modelsonly to compute intrinsic rewards, as commonly done, our method showcases thatthe self-reinforcing cycle between good models and good exploration also opensup another avenue: zero-shot generalization to downstream tasks via model-basedplanning.
Revised: Instead of using models only to compute intrinsic rewards, as commonly done, our method showcases that the self-reinforcing cycle between good models and good exploration also opens up another avenue: zero-shot generalization to downstream tasks via model-based planning.

Original: After the entirely intrinsic task-agnostic exploration phase, our methodsolves challenging downstream tasks such as stacking, flipping, pick & place, andthrowing and generalizes to unseen numbers and arrangements of objects withoutany additional training.
Revised: After the entirely intrinsic task-agnostic exploration phase, our method solves challenging downstream tasks such as stacking, flipping, pick & place, and throwing and generalizes to unseen numbers and arrangements of objects without any additional training.

Original:  Introduction
Revised:  Introduction

Original: Curious free-play has been identified as a driving force in child development, allowing children toefficiently explore their environment and build an understanding of the world [1].
Revised: Curious free-play has been identified as a driving force in child development, allowing children to efficiently explore their environment and build an understanding of the world [1].

Original: Such intrinsicallymotivated exploration schemes are especially attractive in open-ended learning scenarios to guidean agent even without extrinsic tasks and corresponding rewards.
Revised: Such intrinsically motivated exploration schemes are especially attractive in open-ended learning scenarios to guide an agent even without extrinsic tasks and corresponding rewards.

Original: Similar to how children learn, wewant Reinforcement Learning (RL) agents to learn through play and then be able to solve new tasksquickly.
Revised: Similar to how children learn, we want Reinforcement Learning (RL) agents to learn through play and then be able to solve new tasks quickly.

Original: We address both challenges in this work.
Revised: We address both challenges in this work.

Original: Minimization of novelty (or surprise) is a prominent formulation of curiosity, with several psychological studies showcasing the role of novelty in children’s curious exploration [2–4].
Revised: Minimization of novelty (or surprise) is a prominent formulation of curiosity, with several psychological studies showcasing the role of novelty in children’s curious exploration [2–4].

Original: Novelty asan intrinsic reward signal has also been adopted in RL, where agents try to resolve a cognitivedisequilibrium [5, 6].
Revised: Novelty as an intrinsic reward signal has also been adopted in RL, where agents try to resolve a cognitive disequilibrium [5, 6].

Original: However, applying entirely intrinsic task-agnostic exploration to object manipulation scenarios in a sample-efficient manner is an ongoing challenge as the relevant informationlies in the sparse agent-object and object-object interactions.
Revised: However, applying entirely intrinsic task-agnostic exploration to object manipulation scenarios in a sample-efficient manner is an ongoing challenge as the relevant information lies in the sparse agent-object and object-object interactions.

Original: These compositional multi-objectmanipulation environments highlight one of the significant weaknesses of the current novelty-basedintrinsic motivation methods: a novel stimulus alone does not necessarily mean that it contains usefulor generalizable information for an individual [7].
Revised: These compositional multi-object manipulation environments highlight one of the significant weaknesses of the current novelty-based intrinsic motivation methods: a novel stimulus alone does not necessarily mean that it contains useful yields capable models for zero-shot task planning.

Original: Thus, curious exploration needs to attend to asubset of possibilities, as supported by studies in psychology [8, 9].

Original: Analyzing children during free yields capable models for zero-shot task planning.

Original: In an intrinsic phase, the system plans for maximum epistemic uncertainty with an ensemble of GNNs and uses the actively collected data to update the model.
Revised: In an intrinsic phase, the system plans for maximum epistemic uncertainty with an ensemble of GNNs and uses the actively collected data to update the model.

Original: This iterative procedure ultimately reduces the overall epistemic uncertainty of the system, i.e. maximizing information gain.
Revised: This iterative procedure ultimately reduces the overall epistemic uncertainty of the system, i.e. maximizing information gain.

Original: The learned models can be used to solve complex manipulation tasks using online planning.
Revised: The learned models can be used to solve complex manipulation tasks using online planning.

Revised: or generalizable information for an individual [7].

Revised: Thus, curious exploration needs to attend to a subset of possibilities, as supported by studies in psychology [8, 9].

Original:  play accumulated evidence that infants have innate biases and heuristics to help guide their attentiontoward relevant and informative features of the environment [10, 11].
Revised: Analyzing children during free play accumulated evidence that infants have innate biases and heuristics to help guide their attention toward relevant and informative features of the environment [10, 11].

Original: Our goal is to improve curiousexploration in RL by incorporating appropriate inductive biases.
Revised: Our goal is to improve curious exploration in RL by incorporating appropriate inductive biases.

Original: We hypothesize that viewing theworld as a collection of entities and their interactions is such a “useful” inductive bias, which wasalso put forward in Tsividis et al.
Revised: We hypothesize that viewing the world as a collection of entities and their interactions is such a “useful” inductive bias, which was also put forward in Tsividis et al.

Original: However, instead of utilizing explicit theory-based modeling,we exploit a relational inductive bias by using Graph Neural Networks (GNN) as our choice of model[13] and learn everything from interactions.
Revised: However, instead of utilizing explicit theory-based modeling, we exploit a relational inductive bias by using Graph Neural Networks (GNN) as our choice of model [13] and learn everything from interactions.

Original: We postulate that curious exploration via epistemicuncertainty of such a model leads to the collection of valuable data as it inherits this structure.
Revised: We postulate that curious exploration via epistemic uncertainty of such a model leads to the collection of valuable data as it inherits this structure.

Original: This paper shows how online planning methods can efficiently exploit learned models, both for exploration and zero-shot generalization to tasks.
Revised: This paper shows how online planning methods can efficiently exploit learned models, both for exploration and zero-shot generalization to tasks.

Original: Recent advances in general-purpose model predictivecontrol methods have reduced their complexity and show strong performance when good forwardmodels are available [14, 15].
Revised: Recent advances in general-purpose model predictive control methods have reduced their complexity and show strong performance when good forward models are available [14, 15].

Original: We build on this work to create a self-reinforcing cycle betweenlearning good models and good exploration.
Revised: We build on this work to create a self-reinforcing cycle between learning good models and good exploration.

Original: Thanks to their ability to plan for a new task withoutfurther training, models efficiently track the naturally changing exploration targets and performdownstream tasks in a zero-shot generalization manner.
Revised: Thanks to their ability to plan for a new task without further training, models efficiently track the naturally changing exploration targets and perform downstream tasks in a zero-shot generalization manner.

Original: We propose CEE-US: C urious E xploration using E pistemic U ncertainty via S tructured Models thatachieves sample-efficient and interaction-rich exploration in multi-object manipulation environments.
Revised: We propose CEE-US: C urious E xploration using E pistemic U ncertainty via S tructured Models that achieves sample-efficient and interaction-rich exploration in multi-object manipulation environments.

Original: To our knowledge, we are the first to use GNN ensemble disagreement for computing intrinsicmotivation signals in RL.
Revised: To our knowledge, we are the first to use GNN ensemble disagreement for computing intrinsic motivation signals in RL.

Original: A further major contribution of our method is to combine GNN-basedepistemic uncertainty with planning methods in a task-agnostic setting and demonstrate zero-shotgeneralization to challenging object manipulation tasks.
Revised: A further major contribution of our method is to combine GNN-based epistemic uncertainty with planning methods in a task-agnostic setting and demonstrate zero-shot generalization to challenging object manipulation tasks.

Original: We provide a brief comparison to relatedwork in Table 1 and provide more details in Sec.
Revised: We provide a brief comparison to related work in Table 1 and provide more details in Sec.

Original: [Table] [Table]
Revised: [Table] [Table]

Original:  Method
Revised:  Method

Original: We focus on intrinsically motivated learning to prepare for future tasks.
Revised: We focus on intrinsically motivated learning to prepare for future tasks.

Original: Thus, we have a task-agnostic
Revised: Thus, we have a task-agnostic

Original: RL setting without extrinsic rewards or any other information regarding future tasks/goals during theinitial play phase.
Revised: RL setting without extrinsic rewards or any other information regarding future tasks/goals during the initial play phase.

Original: Our approach trains a structured world model to capture the forward dynamics ofthe environment from active exploration data.
Revised: Our approach trains a structured world model to capture the forward dynamics of the environment from active exploration data.

Original: In particular, we are using a Graph Neural Network(GNN) where the nodes correspond to objects.
Revised: In particular, we are using a Graph Neural Network (GNN) where the nodes correspond to objects.

Original: We rely on our model’s epistemic uncertainty to directexploration during free play.
Revised: We rely on our model’s epistemic uncertainty to direct exploration during free play.

Original: Our focus is on compositional multi-object environments, where anactuated agent can independently manipulate different objects.
Revised: Our focus is on compositional multi-object environments, where an actuated agent can independently manipulate different objects.

Original: The terms (actuated) agent and robotare used interchangeably.
Revised: The terms (actuated) agent and robot are used interchangeably.

Original: 2.1 Preliminaries
Revised: 2.1 Preliminaries

Original: In this work, we consider a fully observable Markov Decision Process (MDP) setting, The MDP isgiven by M = ( S , A , P ass ′ , R ass ′ ) , with the continuous state-space S ∈ R n s , the continuous actionspace A ∈ R n a , the transition kernel P ass ′ , and the reward function R ass ′ .
Revised: In this work, we consider a fully observable Markov Decision Process (MDP) setting, The MDP is given by M = ( S , A , P ass ′ , R ass ′ ) , with the continuous state-space S ∈ R n s , the continuous actionspace A ∈ R n a , the transition kernel P ass ′ , and the reward function R ass ′ .

Original: Furthermore, we consideran object-oriented state representation, i.e. the state-space factorizes into the different entities in theenvironment S = S 1 × S 2 · · · × S N × S agent , where N denotes the number of objects.
Revised: Furthermore, we consider an object-oriented state representation, i.e. the state-space factorizes into the different entities in the environment S = S 1 × S 2 · · · × S N × S agent , where N denotes the number of objects.

Original: In the typical RL setting, the agent’s goal is to learn a policy a ∼ π ( · | s ) that maximizes the future(discounted) cumulative reward G t = (cid:80) ∞ k =0 γ k · R t + k , with the discount factor 0 < γ ≤ 1 .
Revised: In the typical RL setting, the agent’s goal is to learn a policy a ∼ π ( · | s ) that maximizes the future (discounted) cumulative reward G t = (cid:80) ∞ k =0 γ k · R t + k , with the discount factor 0 < γ ≤ 1 .

Original: Such apolicy can be learned as a neural network with RL algorithms, or a planning method can be used tomaximize the same quantity.
Revised: Such a policy can be learned as a neural network with RL algorithms, or a planning method can be used to maximize the same quantity.

Original: 2.1.1 Planning and Model Predictive Control
Revised: 2.1.1 Planning and Model Predictive Control

Original: Planning methods use a model f of the transition kernel P ass ′ and a reward function R ass ′ to optimizea sequence of actions on the fly.
Revised: Planning methods use a model f of the transition kernel P ass ′ and a reward function R ass ′ to optimize a sequence of actions on the fly.

Original: The potential advantage of planning methods is that they canoptimize for different reward functions without further adaptation.
Revised: The potential advantage of planning methods is that they can optimize for different reward functions without further adaptation.

Original: For this to be successful, a goodtransition model f needs to be learned, and the reward function needs to be known or discovered.
Revised: For this to be successful, a good transition model f needs to be learned, and the reward function needs to be known or discovered.

Original: More formally, for a fixed planning horizon H , the action sequence a t =
Revised: More formally, for a fixed planning horizon H , the action sequence a t =

Original: ( a t + h ) H − 1 h =0 is optimized tomaximize
Revised: [Equation] a [Equation] =0 is optimized to maximize

Original: [Equation]
Revised: [Equation]

Original: where s t + h are imagined states visited by rolling out the actions using f , which is assumed to bedeterministic.
Revised: where s t + h are imagined states visited by rolling out the actions using f , which is assumed to be deterministic.

Original: For this procedure to optimize the same infinite horizon reward with discounting, γand a value function V ( s t 
Revised: For this procedure to optimize the same infinite horizon reward with discounting, γ and a value function V ( s t + H ) could be added to Eq.

Original: + H ) could be added to Eq. 1, sacrificing the flexibility in exchanging R .
Revised: 1, sacrificing the flexibility in exchanging R .

Original: We use zero-order trajectory optimization to find actions according to Eq.
Revised: We use zero-order trajectory optimization to find actions according to Eq.

Original: 1 and model predictivecontrol (MPC) to convert the open-loop planning policy into a closed-loop policy (re-planningafter every step in the environment).
Revised: 1 and model predictive control (MPC) to convert the open-loop planning policy into a closed-loop policy (re-planning after every step in the environment).

Original: In particular, we use the improved Cross-Entropy Method(iCEM) [14], which was recently proposed as a sample efficient trajectory optimizer.
Revised: In particular, we use the improved Cross-Entropy Method (iCEM) [14], which was recently proposed as a sample efficient trajectory optimizer.

Original: 2.2 World Model with Graph Neural Networks
Revised: 2.2 World Model with Graph Neural Networks

Original: As the forward dynamics f : S ×
Revised: As the forward dynamics f : S ×

Original: A → S (world model) is an integral part of the planningprocedure, it needs to be able to capture the true dynamics well.
Revised: A → S (world model) is an integral part of the planning procedure, it needs to be able to capture the true dynamics well.

Original: We employ an ensemble ofmessage-passing GNNs [13].
Revised: We employ an ensemble of message-passing GNNs [13].

Original: Each object corresponds to a node in the graph, and the node attributes{ s nt ∈ S n | n = 1 , . . .
Revised: Each object corresponds to a node in the graph, and the node attributes { s nt ∈ S n | n = 1 , . . .

Original: , N } are given by an object’s features such as position, orientation, andvelocity at timestep t .
Revised: , N } are given by an object’s features such as position, orientation, and velocity at timestep t .

Original: The state representation of the actuated agent s agent ∈ S agent similarly containsposition and velocity information about the robot.
Revised: The state representation of the actuated agent s agent ∈ S agent similarly contains position and velocity information about the robot.

Original: The agent representation is differentiated from the object nodes as it has a direct cause-and-effectrelationship with the actions.
Revised: The agent representation is differentiated from the object nodes as it has a direct cause-and-effect relationship with the actions.

Original: The agent’s state and the generic action are represented as a globalcontext c .
Revised: The agent’s state and the generic action are represented as a global context c .

Original: The GNN is fully connected, i.e. there is an edge between all nodes.
Revised: The GNN is fully connected, i.e. there is an edge between all nodes.

Original: In GNNs, the nodeupdate function g node models the dynamics of individual entities, and the edge update function g edgecaptures their pairwise 
Revised: In GNNs, the node update function g node models the dynamics of individual entities, and the edge update function g edge captures their pairwise interactions.

Original: [Figure]
Revised: [Figure]

Revised: Along the time axis, 2000 transitions correspond to one training iteration of CEE-US, with 20 rollouts (episodes) of length 100.

Revised: After only 30 iterations of free play, the robot starts lifting objects.

Original: The node attribute computation for the objects is given by:
Revised: The node attribute computation for the objects is given by:

Original: [Equation]
Revised: [Equation]

Original: where [ · , . .
Revised: where [ · , . .

Original: denotes concatenation, e ( i,j ) t is the edge attribute between two neighboring nodes ( i, j ) ,and c is the global context information.
Revised: denotes concatenation, e ( i,j ) t is the edge attribute between two neighboring nodes ( i, j ) , and c is the global context information.

Original: The permutation-invariant aggregation function is given byaggr .
Revised: The permutation-invariant aggregation function is given by aggr .

Original: The g node and g edge are both Multilayer Perceptrons (MLP).
Revised: The g node and g edge are both Multilayer Perceptrons (MLP).

Original: The context update, i.e. the transition of the agent’s state, is computed using the global aggregation ofall edges similar to the global context formalism in Battaglia et al.
Revised: The context update, i.e. the transition of the agent’s state, is computed using the global aggregation of all edges similar to the global context formalism in Battaglia et al.

Original: [Equation]
Revised: [Equation]

Original: where g global denotes the global node MLP.
Revised: where g global denotes the global node MLP.

Original: By providing all information to the prediction of theagent, we ensure an accurate modeling of agent-object interactions and their influence on the agent,which is paramount in object manipulation environments.
Revised: By providing all information to the prediction of the agent, we ensure an accurate modeling of agent-object interactions and their influence on the agent, which is paramount in object manipulation environments.

Original: Using this representation, we found that asingle message passing step is sufficient, thus yielding fast inference times.
Revised: Using this representation, we found that a single message passing step is sufficient, thus yielding fast inference times.

Original: Moreover, to make the model focus on capturing the changes accurately, we let the GNN predict∆ˆ s t +

Revised: Moreover, to make the model focus on capturing the changes accurately, we let the GNN predict ∆ˆ s t +

Original: = ˆ s t +1 − s t , instead of the absolute next state, such that f ( s, a ) =
Revised: = ˆ s t +1 − s t , instead of the absolute next state, such that f ( s, a ) =

Original: s + GNN ( s, a )
Revised: s + GNN ( s, a )

Original: 2.3 Epistemic Uncertainty as Intrinsic Reward
Revised: 2.3 Epistemic Uncertainty as Intrinsic Reward

Original: We train an ensemble of GNNs { ( f θm )
Revised: We train an ensemble of GNNs { ( f θm )

Original: Mm =1 } (Fig. 1), where M denotes the ensemble size.
Revised: Mm =1 } (Fig. 1), where M denotes the ensemble size.

Original: Theepistemic uncertainty, i.e. the uncertainty due to lack of data, can be approximated by the disagreementof the ensemble members’ predictions measured by the trace of the covariance matrix [15]:
Revised: The epistemic uncertainty, i.e. the uncertainty due to lack of data, can be approximated by the disagreement of the ensemble members’ predictions measured by the trace of the covariance matrix [15]:

Original: [Equation]
Revised: [Equation]

Original: Connections to Information Gain If more data becomes available from a region where theensemble members disagree, then we expect the disagreement in this region to shrink, given enoughcapacity of the predictors.
Revised: Connections to Information Gain If more data becomes available from a region where the ensemble members disagree, then we expect the disagreement in this region to shrink, given enough capacity of the predictors.

Original: By guiding the exploration to regions with large model disagreement orhigh epistemic uncertainty, we explicitly aim for maximizing information gain [19–21].
Revised: By guiding the exploration to regions with large model disagreement or high epistemic uncertainty, we explicitly aim for maximizing information gain [19–21].

Original: Interestingly,as we can predict the uncertainty for unvisited states, our approach is more related to predictedinformation gain [22].
Revised: Interestingly, as we can predict the uncertainty for unvisited states, our approach is more related to predicted information gain [22].

Original: 2.4 The CEE-US Algorithm
Revised: 2.4 The CEE-US Algorithm

Original: CEE-US has two different phases: the intrinsic free-play phase (Alg. 1), where learning occurs withactive exploration, and the extrinsic phase (Alg. 2), where we apply the learned model to solvedownstream tasks zero-shot without any additional training.
Revised: CEE-US has two different phases: the intrinsic free-play phase (Alg. 1), where learning occurs with active exploration, and the extrinsic phase (Alg. 2), where we apply the learned model to solve downstream tasks zero-shot without any additional training.

Original: Intrinsic Phase
Revised: Intrinsic Phase

Original: In each iteration, the zero-order trajectory optimizer is used to plan for actionsequences with high cumulative epistemic uncertainty in the imagined trajectories of the world model.
Revised: In each iteration, the zero-order trajectory optimizer is used to plan for action sequences with high cumulative epistemic uncertainty in the imagined trajectories of the world model.

Original: The collected rollouts are added to the buffer, and the ensemble members of the world model are trained on the observed state transitions ( s t , a t , s t [Equation] 1 ) for a fixed number of epochs minimizing theloss function:
Revised: The collected rollouts are added to the buffer, and the ensemble members of the world model are trained on the observed state transitions ( s t , a t , s t +1 ) for a fixed number of epochs minimizing the loss function:

Original: [Equation]
Revised: [Equation]

Original: (7) for each ensemble member m with independently sampled mini-batches, where ∆ s t +1 =
Revised: (7) for each ensemble member m with independently sampled mini-batches, where ∆ s t +1 =

Original: s t +1 − s tdenotes the actual change in the next state observation.
Revised: s t +1 − s t denotes the actual change in the next state observation.

Original: The total loss function is given by:
Revised: The total loss function is given by:

Original: L = (cid:80) Mi =

Revised: L = (cid:80) Mi =

Original: L m .
Revised: L m .

Original: Afterwards, the process of data collection and training is repeated.
Revised: Afterwards, the process of data collection and training is repeated.

Original: The pseudocode for the intrinsic phase is provided in Alg.
Revised: The pseudocode for the intrinsic phase is provided in Alg.

Original: Lines 3–9 correspond to one training iteration.
Revised: Lines 3–9 correspond to one training iteration.

Original: [Table] [Table]
Revised: [Table] [Table]

Original: Extrinsic Phase
Revised: Extrinsic Phase

Original: The learned world model is used to perform planning for several downstreamtasks, which we assume to be given in terms of reward functions R task ( s, a, s ′ ) .
Revised: The learned world model is used to perform planning for several downstream tasks, which we assume to be given in terms of reward functions R task ( s, a, s ′ ) .

Original: The pseudocode ispresented in Alg.2.
Revised: The pseudocode is presented in Alg.

Original: Here, no further adaptation to the model is performed, although our methodcould be extended to incorporate a potential fine-tuning procedure.
Revised: Here, no further adaptation to the model is performed, although our method could be extended to incorporate a potential fine-tuning procedure.

Original:  Experiments
Revised:  Experiments

Original: In our empirical evaluation, we analyze the performance of CEE-US on two different object manipulation environments to answer the questions: How much does the structural inductive bias introducedby GNNs help model learning and control?
Revised: In our empirical evaluation, we analyze the performance of CEE-US on two different object manipulation environments to answer the questions: How much does the structural inductive bias introduced by GNNs help model learning and control?

Original: Does the free-play phase create rich interaction data thathelps downstream task performance?
Revised: Does the free-play phase create rich interaction data that helps downstream task performance?

Original: Can we solve challenging manipulation tasks in a zero-shotmanner?
Revised: Can we solve challenging manipulation tasks in a zero-shot manner?

Original: The two environments we consider are:
Revised: The two environments we consider are:

Original: Playground An actuated spherical agent can slide in the x-y directions and push four different object types (light cube, heavy cube, pyramid, and cylinder) along the x-y directions and also rotate them around the z-axis (Fig. 
Revised: Playground An actuated spherical agent can slide in the x-y directions and push four different object types (light cube, heavy cube, pyramid, and cylinder) along the x-y directions and also rotate them around the z-axis (Fig. 3).

Original: Each object has 3 Degrees of Freedom (DoF).
Revised: Each object has 3 Degrees of Freedom (DoF).

Original: Object types are uniquely identified by their color
Revised: Object types are uniquely identified by their color.

Original: Fetch Pick & Place Construction This is an extension of the Fetch Pick & Place environment [23]to more cubes [24] (Fig. 2).
Revised: Fetch Pick & Place Construction This is an extension of the Fetch Pick & Place environment [23] to more cubes [24] (Fig. 2).

Original: A 7-DoF robot arm is used to manipulate blocks (each with 6 DoF).
Revised: A 7-DoF robot arm is used to manipulate blocks (each with 6 DoF).

Original: Theactions a ∈ R 4 control the gripper movement in Cartesian coordinates and the gripper opening/closing.
Revised: The actions a ∈ R 4 control the gripper movement in Cartesian coordinates and the gripper opening/closing.

Original: The robot state s agent ∈ R 10 contains the end-effector position and velocity and the gripper-state(open/close) and gripper-velocity.
Revised: The robot state s agent ∈ R 10 contains the end-effector position and velocity and the gripper-state (open/close) and gripper-velocity.

Original: Each object’s state s i ∈
Revised: Each object’s state s i ∈

Original: R 12 is given by its position, orientation(in Euler angles), and linear and angular velocities.
Revised: R 12 is given by its position, orientation (in Euler angles), and linear and angular velocities.

Original: We replaced the table in front of the robot witha large plane so that objects cannot fall off during free play.
Revised: We replaced the table in front of the robot with a large plane so that objects cannot fall off during free play.

Original: However, they can still be pushed orthrown outside of the manipulability range of the robot.
Revised: However, they can still be pushed or thrown outside of the manipulability range of the robot.

Original: Originally, each object state contained theobject’s position relative to the gripper.
Revised: Originally, each object state contained the object’s position relative to the gripper.

Original: We remove this privileged information from the objects’ state,as it already introduces a relational bias in the raw state representation.
Revised: We remove this privileged information from the objects’ state, as it already introduces a relational bias in the raw state representation.

Original: 3.1 Structured vs. Unstructured World Models
Revised: 3.1 Structured vs. Unstructured World Models

Original: We analyze how curious exploration based on minimizing epistemic uncertainty performs when usingensembles of GNNs and ensembles of MLPs (Fig. 3).
Revised: We analyze how curious exploration based on minimizing epistemic uncertainty performs when using ensembles of GNNs and ensembles of MLPs (Fig. 3).

Original: MLPs, with their fully connected layers andmonolithic input representations, do not offer isolation of information and incorporate no explicitrelational inductive bias [13], thus constituting a good baseline for our structured world models.
Revised: MLPs, with their fully connected layers and monolithic input representations, do not offer isolation of information and incorporate no explicit relational inductive bias [13], thus constituting a good baseline for our structured world models.

Original: Results in P LAYGROUND show that the GNN ensemble leads much faster to interaction-rich data
Revised: Results in P LAYGROUND show that the GNN ensemble leads much faster to interaction-rich data

Original: [Figure] [Figure]
Revised: [Figure] [Figure]

Original: than its MLP counterpart.
Revised: than its MLP counterpart.

Original: As visualized in Fig.
Revised: As visualized in Fig.

Original: 3, the uncertainty produced by the GNN ensembleis already localized around objects after 25 training iterations, leading to targeted agent-object andobject-object interactions.
Revised: [Figure] the uncertainty produced by the GNN ensemble is already localized around objects after 25 training iterations, leading to targeted agent-object and object-object interactions.

Original: For the MLP, it takes more than 100 iterations to start generating usefuluncertainty estimates and therefore object-agent interactions.
Revised: For the MLP, it takes more than 100 iterations to start generating useful uncertainty estimates and therefore object-agent interactions.

Original: Also, note the pronounced uncertaintyat the walls for the MLP.
Revised: Also, note the pronounced uncertainty at the walls for the MLP.

Original: A more fine-grained analysis of the interaction times of CEE-US comparedto the baselines will be provided below. 
Revised: A more fine-grained analysis of the interaction times of CEE-US compared to the baselines will be provided below. In Suppl.

Revised: D, we also show how the resulting multi-step dynamics predictions with structured world models are more accurate compared to the MLPs.

Original: 3.2 Interaction-Richness of Generated Exploration Data
Revised: 3.2 Interaction-Richness of Generated Exploration Data

Original: We perform an analysis of the data generated during intrinsic free-play.
Revised: We perform an analysis of the data generated during intrinsic free-play.

Original: We compare the performanceof our method to the unstructured MLP + iCEM variant (CEE-US without GNNs) as well as otherintrinsic motivation baselines:
Revised: We compare the performance of our method to the unstructured MLP + iCEM variant (CEE-US without GNNs) as well as other intrinsic motivation baselines:

Original: Disagreement [16]
Revised: Disagreement [16]

Original: The one-step disagreement of an MLP ensemble is used as an intrinsic rewardto train an exploration policy.
Revised: The one-step disagreement of an MLP ensemble is used as an intrinsic reward to train an exploration policy.

Original: Random Network Distillation [17] A predictor network, corresponding to the forward model,tries to match the output of a target network with random weights.
Revised: Random Network Distillation [17] A predictor network, corresponding to the forward model, tries to match the output of a target network with random weights.

Original: The discrepancy between the twonetworks is used as an intrinsic reward (a type of state visitation count for continuous domains).
Revised: The discrepancy between the two networks is used as an intrinsic reward (a type of state visitation count for continuous domains).

Original: ICM [6] The intrinsic reward is defined as the error between an MLP forward model’s next stateprediction and the actual next state.
Revised: ICM [6] The intrinsic reward is defined as the error between an MLP forward model’s next state prediction and the actual next state.

Original: As this method needs access to the true next state, the intrinsicreward can only be computed retrospectively.
Revised: As this method needs access to the true next state, the intrinsic reward can only be computed retrospectively.

Original: The intrinsic reward is used to train an exploration policy in these methods.
Revised: The intrinsic reward is used to train an exploration policy in these methods.

Original: We use the implementationfrom Laskin et al.
Revised: We use the implementation from Laskin et al.

Original: [25] that uses DDPG [26].
Revised: [25] that uses DDPG [26].

Original: In Fig.
Revised: In Fig.

Original: 4, we present different metrics quantifyingtheamount and type of interactions that occur during the free play in the P LAYGROUND environment.
Revised: 4, we present different metrics quantifying the [Figure] MLP

Original: [Figure]
Revised: [Figure]

Revised: amount and type of interactions that occur during the free play in the P LAYGROUND environment.

Original: It is eminent that CEE-US is only shortly interacting with the walls and then quickly interacts withone and two objects simultaneously.
Revised: It is eminent that CEE-US is only shortly interacting with the walls and then quickly interacts with one and two objects simultaneously.

Original: The next best method is the ablation of our method using MLPmodels needing about ten times more interactions, as explained in Fig.
Revised: The next best method is the ablation of our method using MLP models needing about ten times more interactions, as explained in Fig.

Original: The baselines are mainlymoving in free space and interacting with the walls.
Revised: The baselines are mainly moving in free space and interacting with the walls.

Original: RND starts interacting with the objects in at least10% of the times-steps.
Revised: RND starts interacting with the objects in at least 10% of the times-steps.

Original: In the C ONSTRUCTION environment, the situation is even more drastic, as shown in Fig.
Revised: In the C ONSTRUCTION environment, the situation is even more drastic, as shown in Fig.

Original: CEE-USstarts repeatedly moving one object after only eight iterations (each 2000 environment steps), picksup objects after around 30 iterations, and continues with throwing and flipping objects and movingmultiple objects frequently after 50 iterations.
Revised: CEE-US starts repeatedly moving one object after only eight iterations (each 2000 environment steps), picks up objects after around 30 iterations, and continues with throwing and flipping objects and moving multiple objects frequently after 50 iterations.

Original: In Fig.
Revised: In Fig.

Original: 2 exemplary behaviors are shown in termsof snapshots at different stages of learning.
Revised: 2 exemplary behaviors are shown in terms of snapshots at different stages of learning.

Original: We believe this sample efficiency is remarkable.
Revised: We believe this sample efficiency is remarkable.

Original: Theablation of our method with MLP models is also able to engage in interesting and diverse interactions,but at a much slower rate (Fig. 5).
Revised: The ablation of our method with MLP models is also able to engage in interesting and diverse interactions, but at a much slower rate (Fig. 5).

Original: The policy learning baselines rarely interact with objects, evenafter 300 iterations with 20 episodes each.
Revised: The policy learning baselines rarely interact with objects, even after 300 iterations with 20 episodes each.

Original: 3.3 Zero-Shot Generalization to Downstream Tasks
Revised: 3.3 Zero-Shot Generalization to Downstream Tasks

Original: [Figure] We demonstrate that after the free-play phase, the learned models can be used for zero-shot solving ofcomplex downstream tasks without the necessity to generate new data or to perform further training.
Revised: We demonstrate that after the free-play phase, the learned models can be used for zero-shot solving of complex downstream tasks without the necessity to generate new data or to perform further training.

Original: Furthermore, thanks to the combinatorial generalization capabilities of GNNs, the model can be usedwith a different number of objects in the environment than seen during free play.
Revised: Furthermore, thanks to the combinatorial generalization capabilities of GNNs, the model can be used with a different number of objects in the environment than seen during free play.

Original: The only baseline weconsider for the zero-shot generalization performance is MLP + iCEM, as none of the other methodsdiscussed in the previous section can solve downstream tasks without additional training.
Revised: The only baseline we consider for the zero-shot generalization performance is MLP + iCEM, as none of the other methods discussed in the previous section can solve downstream tasks without additional training.

Original: In P LAYGROUND , we consider the task of bringing all objects to a single target location, as shownin Fig.
Revised: In P LAYGROUND , we consider the task of bringing all objects to a single target location, as shown in Fig.

Original: The reward function is the sum of negative distances of all objects to the target outside athreshold distance.
Revised: The reward function is the sum of negative distances of all objects to the target outside a threshold distance.

Original: The target location is sampled randomly at the beginning of each episode.
Revised: The target location is sampled randomly at the beginning of each episode.

Original: Thesuccess rates for different numbers of objects are shown in Table 2.
Revised: The success rates for different numbers of objects are shown in Table 2.

Original: We define success rate in themulti-object setup as the fraction of objects solved relative to the total number of objects spawned inthe environment.
Revised: We define success rate in the multi-object setup as the fraction of objects solved relative to the total number of objects spawned in the environment.

Original: After 100 training iterations of free play, CEE-US already achieves a 70% successrate, whereas the MLP version only reaches 8%.
Revised: After 100 training iterations of free play, CEE-US already achieves a 70% success rate, whereas the MLP version only reaches 8%.

Original: After 250 iterations, both methods are on par.
Revised: After 250 iterations, both methods are on par.

Original: Since
Revised: Since

Original:  CEE-US can also deal with a variable input dimension, we also consider the task with more or lessobjects with randomly sampled types.[Figure]
Revised: [Figure] [Figure] CEE-US can also deal with a variable input dimension, we also consider the task with more or less objects with randomly sampled types.

Original: [Figure]

Original: [Table]

Original: [Figure]
Revised: [Figure]

Original: In C ONSTRUCTION , several challenging manipulation tasks need to be solved: pick & place,stacking, throwing, and flipping, as shown in
Revised: In C ONSTRUCTION , several challenging manipulation tasks need to be solved: pick & place, stacking, throwing, and flipping, as shown in

Original: Fig.
Revised: Fig.

Original: The reward functions are detailed in
Revised: The reward functions are detailed in

Original: Suppl.
Revised: Suppl.

Original: C.3.1.
Revised: C.3.1.

Original: In Fig.
Revised: In Fig.

Original: 7, we present the success rates in each of the tasks for CEE-US andthe MLP-based planning baseline MLP + iCEM,where applicable.
Revised: 7, we present the success rates in each of the tasks for CEE-US and the MLP-based planning baseline MLP + iCEM, where applicable.

Original: We find remarkable successrates across the board for CEE-US, even in challenging tasks such as stacking and throwing.
Revised: We find remarkable success rates across the board for CEE-US, even in challenging tasks such as stacking and throwing.

Original: Forstacking, we report a success rate of 1 if thetower with all the objects in the environment isstacked and 0 otherwise.
Revised: For stacking, we report a success rate of 1 if the tower with all the objects in the environment is stacked and 0 otherwise.

Original: For the multi-towertask, denoted by 2 + 2 , we consider success 1only when two towers are stacked.
Revised: For the multi-tower task, denoted by 2 + 2 , we consider successonly when two towers are stacked.

Original: We recommend visiting our website 2 for videos of these tasks. 
Revised: We recommend visiting our website 2 for videos of these tasks. In Suppl.

Revised: E, we present additional experiments combining RND as intrinsic reward with model-based planning in C ONSTRUCTION , where we look at both structured and unstructured world models.

Revised: The results showcase the benefits of using the model’s own epistemic uncertainty estimate to guide exploration, as it is the case with ensemble disagreement, leading to more accurate and more robust dynamics models and better zero-shot downstream task generalization.

Original: 3.4 Offline Learning of Downstream Tasks from Exploratory Data
Revised: 3.4 Offline Learning of Downstream Tasks from Exploratory Data

Original: This experiment investigates the quality of the data collected by the different methods in the free-playphase for solving downstream tasks via offline RL.
Revised: This experiment investigates the quality of the data collected by the different methods in the free-play phase for solving downstream tasks via offline RL.

Original: This allows us to compare to other baselines thatdo not support zero-shot generalization.
Revised: This allows us to compare to other baselines that do not support zero-shot generalization.

Original: Additionally, it is one way of obtaining task-specific policies.
Revised: Additionally, it is one way of obtaining task-specific policies.

Original: In offline RL, a control policy is learned from a fixed dataset, where we repurpose the explorationdata.
Revised: In offline RL, a control policy is learned from a fixed dataset, where we repurpose the exploration data.

Original: The results are shown in Table 3.
Revised: The results are shown in Table 3.

Original: Offline policy learning is performed via CQL [27] 3 andtraining details can be found in Suppl.
Revised: Offline policy learning is performed via CQL [27] 3 and training details can be found in Suppl.

Original: C.5.
Revised: C.5.

Original: To learn different downstream tasks from the same dataset,reward relabeling and hindsight experience replay [29] are used.
Revised: To learn different downstream tasks from the same dataset, reward relabeling and hindsight experience replay [29] are used.

Original: In accordance with the interaction metrics during the intrinsic phase (Figs. 4,5), CEE-US and
Revised: In accordance with the interaction metrics during the intrinsic phase (Figs. 4,5) [Table] CEE-US and

Original: MLP + iCEM achieve the highest performance in all four tasks, supporting the benefit of optimizingfor long-horizon future novelty.
Revised: MLP + iCEM achieve the highest performance in all four tasks, supporting the benefit of optimizing for long-horizon future novelty.

Original: Especially in the object manipulation tasks, data from CEE-US showsa clear advantage.
Revised: Especially in the object manipulation tasks, data from CEE-US shows a clear advantage.

Original: Tasks with more objects could not be solved with the amount of data collected.
Revised: Tasks with more objects could not be solved with the amount of data collected

Original: [Table]
Revised: [Table]

Revised: [Table]

Original:  Related Work
Revised:  Related Work

Original: Intrinsic motivation in RL Prediction error [5, 6, 30], novelty and Bayesian surprise [20, 31, 32],learning progress [5, 31, 33], empowerment [34, 35] and count-based metrics [17, 36] are amongthe popular intrinsic reward signal definitions used in RL.
Revised: Intrinsic motivation in RL Prediction error [5, 6, 30], novelty and Bayesian surprise [20, 31, 32], learning progress [5, 31, 33], empowerment [34, 35] and count-based metrics [17, 36] are among the popular intrinsic reward signal definitions used in RL.

Original: These intrinsic rewards are either usedto aid exploration in challenging tasks with sparse task rewards or in a task-agnostic setup wherethey are the only rewards.
Revised: These intrinsic rewards are either used to aid exploration in challenging tasks with sparse task rewards or in a task-agnostic setup where they are the only rewards.

Original: Algorithms using the task-agnostic setting follow two routes to solvedownstream tasks: (i) they re-label the collected data during intrinsic exploration with the downstreamtask reward and perform offline RL [18, 36, 37] or (ii) they use snapshots of the exploration policyfor bootstrapping  and fine-tune them on downstream tasks [25, 38].
Revised: Algorithms using the task-agnostic setting follow two routes to solve downstream tasks: (i) they re-label the collected data during intrinsic exploration with the downstream task reward and perform offline RL [18, 36, 37] or (ii) they use snapshots of the exploration policy for bootstrapping We are and fine-tune them on downstream tasks [25, 38].

Original: The first variant does not requireadditional interaction with the environment, however, additional training of a new policy is stillnecessary.
Revised: The first variant does not require additional interaction with the environment, however, additional training of a new policy is still necessary.

Original: In addition, offline RL struggles with the distribution shift, while we show that a goodworld model can generalize better.
Revised: In addition, offline RL struggles with the distribution shift, while we show that a good world model can generalize better.

Original: The second variant suffers from another inherent issue: theemerging behaviors of the exploration policy are lost during training.
Revised: The second variant suffers from another inherent issue: the emerging behaviors of the exploration policy are lost during training.

Original: Although Groth et al.
Revised: Although Groth et al.

Original: [38]address this problem by snapshotting and performing hierarchical RL, this is more of a bandaidsolution.
Revised: [38] address this problem by snapshotting and performing hierarchical RL, this is more of a bandaid solution.

Original: Many intrinsically motivated learning systems use a goal-achieving setting with predefined goal spaces [31, 33, 39, 40] or auxiliary tasks [41].
Revised: Many intrinsically motivated learning systems use a goal-achieving setting with predefined goal spaces [31, 33, 39, 40] or auxiliary tasks [41].

Original: The differences in the existing paradigms arehighlighted in Table 1.
Revised: The differences in the existing paradigms are highlighted in Table 1.

Revised: Methods such as ICM [6] rely on retrospective intrinsic motivation (here prediction error), i.e. the agent has to already be in the next state to assess how novel that state is.

Revised: In order to accommodate planning for multi-step intrinsic motivation signals into the future, we need a way to predict them.

Revised: Some methods do so by learning a loss model to predict the prediction error at future time steps [42].

Original:  Plan2Explore [18] shares similarities to our approach in the task-agnosticexploration phase; however they do not capture the directed information-seeking behavior needed forsample efficiency in object manipulation tasks.
Revised: Another line of research, including our method, uses the disagreement of an ensemble of world models to estimate the predictive information gain as the intrinsic motivation measure [16, 18]. Plan2Explore [18] shares similarities to our approach in the task-agnostic exploration phase, using multi-step ensemble disagreement as intrinsic reward. It works with a latent dynamics prediction model and has been applied to domains with image observations. However, Plan2Explore does not use structured world models and lacks mechanisms to achieve combinatorial generalization that is beneficial for sample-efficient exploration in object manipulation tasks.

Original: Relational Networks Several works showcase improved dynamics prediction performance inenvironments with interacting entities using structured world models [42–44].
Revised: Relational Networks Several works showcase improved dynamics prediction performance in environments with interacting entities using structured world models [43–45].

Original: Kipf et al.
Revised: Kipf et al.

Original: [42] uses a GNN to learn the latent transition dynamics in simple manipulation tasks with 2D shapes and 3Dblocks from images.
Revised: [43] uses a GNN to learn the latent transition dynamics in simple manipulation tasks with 2D shapes and 3D blocks from images.

Original: However, they use an object-factorized action space and do not tackle explorationbut rely on an external dataset.
Revised: However, they use an object-factorized action space and do not tackle exploration but rely on an external dataset.

Original: Sanchez-Gonzalez et al.
Revised: Sanchez-Gonzalez et al.

Original: [43] applies GNNs to learn the dynamics ofphysical bodies, where the entities correspond to joints, also without active exploration.
Revised: [44] applies GNNs to learn the dynamics of physical bodies, where the entities correspond to joints, also without active exploration.

Original: Driess et al.[45] uses GNNs and NERFs to obtain object-centric representations from images and RRT planners. 
Revised: Driess et al. [46] uses GNNs and NERFs to obtain object-centric representations from images and RRT planners. Biza et al.

Revised: [47] also utilizes the combinatorial generalization of structured world models to achieve zero-shot task generalization.

Revised: However, the world model is learned on an offline expert dataset, thus sidestepping the exploration challenge, and with parameterized high-level actions.

Original: Outside the model-based paradigm, Li et al.
Revised: Outside the model-based paradigm, Li et al.

Original: [24] achieves block stacking in C ONSTRUCTION using a
Revised: [24] achieves block stacking in C ONSTRUCTION using a

Original: GNN policy, attention and tailored learning curriculum.
Revised: GNN policy, attention and tailored learning curriculum.

Original: In contrast, we use GNNs for computingintrinsic rewards and online planning.
Revised: In contrast, we use GNNs for computing intrinsic rewards and online planning.

Original: Curiosity in Object-Centric RL Watters et al.
Revised: Curiosity in Object-Centric RL Watters et al.

Original: [46] deploy curious exploration with retrospectiveprediction error as an intrinsic reward in an object-centric setting with image inputs and a pre-trainedvision model for object discovery but without object-object interactions.
Revised: [48] deploy curious exploration with retrospective prediction error as an intrinsic reward in an object-centric setting with image inputs and a pre-trained vision model for object discovery but without object-object interactions.

Original: In Seitzer et al.
Revised: In Seitzer et al.

Original: [47], anobject-centric causal action-influence is used as an intrinsic reward to improve sample efficiency.
Revised: [49], an object-centric causal action-influence is used as an intrinsic reward to improve sample efficiency.

Original:  Discussion
Revised:  Discussion

Original: In this work, we present CEE-US that combines the learning of GNNs as structured world modelswith curiosity-driven, planning-based exploration.
Revised: In this work, we present CEE-US that combines the learning of GNNs as structured world models with curiosity-driven, planning-based exploration.

Original: We tackle a significant challenge that is not oftenaddressed in existing intrinsically motivated RL works: guiding exploration towards potentially usefulcomponents of the environment.
Revised: We tackle a significant challenge that is not often addressed in existing intrinsically motivated RL works: guiding exploration towards potentially useful components of the environment.

Original: By approximating information gain via world models injected withrelational inductive biases, CEE-US focuses on the interactions between entities in the environment.
Revised: By approximating information gain via world models injected with relational inductive biases, CEE-US focuses on the interactions between entities in the environment.

Original: Maximizing future information gain via multistep forward-planning enables CEE-US to discoverinteraction-rich behaviors more efficiently than the exploration-policy-based baselines, as well asplanning-based paradigms without structured world models.
Revised: Maximizing future information gain via multistep forward-planning enables CEE-US to discover interaction-rich behaviors more efficiently than the exploration-policy-based baselines, as well as planning-based paradigms without structured world models.

Original: After the intrinsic free-play phase, we use the learned GNNs to solve downstream tasks withmodel-based planning and zero-shot without any additional training.
Revised: After the intrinsic free-play phase, we use the learned GNNs to solve downstream tasks with model-based planning and zero-shot without any additional training.

Original: CEE-US can achieve zeroshot generalization without the need for an additional policy learning step, a strength of CEE-USthat is missing from all the baselines.
Revised: CEE-US can achieve zeroshot generalization without the need for an additional policy learning step, a strength of CEE-US that is missing from all the baselines.

Original: We show that even non-structured MLPs achieve zero-shotgeneralization on some tasks, albeit with lower success rates.
Revised: We show that even non-structured MLPs achieve zero-shot generalization on some tasks, albeit with lower success rates.

Original: This indicates that there are benefitsto learning good world models and utilizing them for control even without any further structuralbiases.
Revised: This indicates that there are benefits to learning good world models and utilizing them for control even without any further structural biases.

Original: This also ties into knowledge-based intrinsic motivation, as the experience gathered duringfree play is distilled and stored in the learned world models instead of being dismissed.
Revised: This also ties into knowledge-based intrinsic motivation, as the experience gathered during free play is distilled and stored in the learned world models instead of being dismissed.

Original: It is importantto note that using model-based planning during the intrinsic phase does not restrict us to model-basedapproaches in the extrinsic phase.
Revised: It is important to note that using model-based planning during the intrinsic phase does not restrict us to model-based approaches in the extrinsic phase.

Original: The learned model could also be used to extract a task policy laterusing model-based policy optimization [48] or to perform offline RL on the data generated duringfree play, which we demonstrate in Sec. 3.4.
Revised: The learned model could also be used to extract a task policy later using model-based policy optimization [50] or to perform offline RL on the data generated during free play, which we demonstrate in Sec. 3.4.

Original: Despite the sample efficiency, there are some limitations to model-based online planning.
Revised: Despite the sample efficiency, there are some limitations to model-based online planning.

Original: Thecomplexity of behaviors discovered in free play is upper-bounded by the finite planning horizon.
Revised: The complexity of behaviors discovered in free play is upper-bounded by the finite planning horizon.

Original: The same also applies to the extrinsic phase, where solving tasks like throwing objects or solvingmultistep manipulation tasks require longer planning horizons without extensive reward shaping.
Revised: The same also applies to the extrinsic phase, where solving tasks like throwing objects or solving multistep manipulation tasks require longer planning horizons without extensive reward shaping.

Original: Although CEE-US currently uses proprioceptive state information, we do not assume access to anyprivileged information such that there are no fundamental limitations prohibiting us from applying itto real robots.
Revised: Although CEE-US currently uses proprioceptive state information, we do not assume access to any privileged information such that there are no fundamental limitations prohibiting us from applying it to real robots.

Original: Following Kipf et al.
Revised: Following Kipf et al.

Original: [42] and Watters et al.
Revised: [43] and Watters et al.

Original: [46], CEE-US can also be extended todeal with image inputs and use methods like SCALOR [49] to extract object-centric unsupervisedrepresentations.
Revised: [48], CEE-US can also be extended to deal with image inputs and use methods like SCALOR [51] to extract object-centric unsupervised representations.

Original: The differentiation between agent and objects can be identified unsupervised [50].
Revised: The differentiation between agent and objects can be identified unsupervised [52].

Original: The demonstrated sample efficiency in the unsupervised learning of capable world models opensnew avenues for learning directly on real hardware.
Revised: The demonstrated sample efficiency in the unsupervised learning of capable world models opens new avenues for learning directly on real hardware.

Original: To put the training time into context, 45 trainingiterations correspond to 1h of interaction time.
Revised: To put the training time into context, 45 training iterations correspond to 1h of interaction time.

Original: So our downstream performance in C ONSTRUCTIONwas achieved after about 6h of free-play from scratch.
Revised: So our downstream performance in C ONSTRUCTION was achieved after about 6.5h of free-play from scratch.

Revised: Acknowledgments and Disclosure of Funding

Original: Checklist
Revised: The authors thank Arash Tavakoli and Pavel Kolev for helpful discussions and Andrii Zadaianchuk, Christian Gumbsch and Nico Gürtler for their help reviewing the manuscript. The authors thank the International Max Planck Research School for Intelligent Systems (IMPRS-IS) for supporting Cansu Sancaktar and Sebastian Blaes.

Original: For all authors...
Revised: Georg Martius is a member of the Machine Learning Cluster of Excellence, EXC number 2064/1 – Project number 390727645. We acknowledge the financial support from the German Federal Ministry of Education and Research (BMBF) through the Tübingen AI Center (FKZ: 01IS18039B). This work was supported by the Volkswagen Stiftung (No 98 571).

Original:  (a) Do the main claims made in the abstract and introduction accurately reflect the paper’scontributions and scope?
Revised: For all authors... (a) Do the main claims made in the abstract and introduction accurately reflect the paper’s contributions and scope?

Original: [Yes] (b) Did you describe the limitations of your work?
Revised: [Yes] (b) Did you describe the limitations of your work?

Original: [Yes] We discuss the limitations of ourwork in Sec. 5.(c)
Revised: [Yes] We discuss the limitations of our work in Sec. 5.

Original:  Did you discuss any potential negative societal impacts of your work?
Revised: (c) Did you discuss any potential negative societal impacts of your work?

Original: [N/A] Webeliever there are no negative societal impacts of our work.(d)
Revised: [N/A] We believer there are no negative societal impacts of our work.

Original:  Have you read the ethics review guidelines and ensured that your paper conforms to them?
Revised: (d) Have you read the ethics review guidelines and ensured that your paper conforms to them?

Original: [Yes] 
Revised: [Yes] 2.

Original: If you are including theoretical results... (a) Did you state the full set of assumptions of all theoretical results?
Revised: If you are including theoretical results... (a) Did you state the full set of assumptions of all theoretical results?

Original: [N/A](b) Did you include complete proofs of all theoretical results?
Revised: [N/A] (b) Did you include complete proofs of all theoretical results?

Original: [N/A] 
Revised: [N/A] 3.

Original: If you ran experiments...
Revised: If you ran experiments...

Original: (a) Did you include the code, data, and instructions needed to reproduce the main exper- [Table] imental results (either in the supplemental material or as a URL)?
Revised: (a) Did you include the code, data, and instructions needed to reproduce the main experimental results (either in the supplemental material or as a URL)?

Original: [Yes] We provide the algorithmic details of our method in Sec.
Revised: [Yes] We provide the algorithmic details of our method in Sec.

Original: 2.4 and further details in supplementarymaterial.
Revised: 2.4 and further details in supplementary material.

Original: The code will also be published later.
Revised: The code is published on the project website.

Original: (b) Did you specify all the training details (e.g., data splits, hyperparameters, how theywere chosen)?
Revised: (b) Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)?

Original: [Yes] See Suppl.
Revised: [Yes] See Suppl.

Original: C[Table] (c) Did you report error bars (e.g., with respect to the random seed after running experiments multiple times)?
Revised: C (c) Did you report error bars (e.g., with respect to the random seed after running experiments multiple times)?

Original: [Yes](d) Did you include the total amount of compute and the type of resources used (e.g., typeof GPUs, internal cluster, or cloud provider)?
Revised: [Yes] (d) Did you include the total amount of compute and the type of resources used (e.g., type of GPUs, internal cluster, or cloud provider)?

Original: [Yes] Further details will be put in the supplementary.
Revised: [Yes] Further details will be put in the supplementary.

Original: [Table] If you are using existing assets (e.g., code, data, models) or curating/releasing new assets...[Table] (a) If your work uses existing assets, did you cite the creators?
Revised: If you are using existing assets (e.g., code, data, models) or curating/releasing new assets... (a) If your work uses existing assets, did you cite the creators?

Original: [Yes] (b) Did you mention the license of the assets?
Revised: [Yes] (b) Did you mention the license of the assets?

Original: [N/A] (c) Did you include any new assets either in the supplemental material or as a URL?
Revised: [N/A] (c) Did you include any new assets either in the supplemental material or as a URL?

Original: [Yes]Code and videos will be published later.
Revised: [Yes] Code and videos will be published later.

Original: [Table] (d) Did you discuss whether and how consent was obtained from people whose data you’re[Table] using/curating?
Revised: (d) Did you discuss whether and how consent was obtained from people whose data you’re using/curating?

Original: [N/A][Table] (e) Did you discuss whether the data you are using/curating contains personally identifiableinformation or offensive content?
Revised: [N/A] (e) Did you discuss whether the data you are using/curating contains personally identifiable information or offensive content?

Original: [N/A]5.
Revised: [N/A] 5.

Original: If you used crowdsourcing or conducted research with human subjects... (a) Did you include the full text of instructions given to participants and screenshots, ifapplicable?
Revised: If you used crowdsourcing or conducted research with human subjects... (a) Did you include the full text of instructions given to participants and screenshots, if applicable?

Original: [N/A](b)
Revised: [N/A]

Original:  Did you describe any potential participant risks, with links to Institutional Review[Table] Board (IRB) approvals, if applicable?
Revised: (b) Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, if applicable?

Original: [N/A][Table] (c) Did you include the estimated hourly wage paid to participants and the total amountspent on participant compensation?
Revised: [N/A] (c) Did you include the estimated hourly wage paid to participants and the total amount spent on participant compensation?

Original: [N/A]
Revised: [N/A]

Original: [Table]