Title: Regularity as Intrinsic Reward for Free Play
Venue: NeurIPS 2023
Accepted: yes
Keywords: Intrinsic Motivation

Reinforcement Learning

Model-based Planning

Regularity

Manipulation

Zero-shot Generalization

Unsupervised Exploration
Abstract: We propose regularity as a novel reward signal for intrinsically-motivated reinforcement learning. Taking inspiration from child development, we postulate that striving for structure and order helps guide exploration towards a subspace of tasks that are not favored by naive uncertainty-based intrinsic rewards. Our generalized formulation of Regularity as Intrinsic Reward (RaIR) allows us to operationalize it within model-based reinforcement learning. In a synthetic environment, we showcase the plethora of structured patterns that can emerge from pursuing this regularity objective. We also demonstrate the strength of our method in a multi-object robotic manipulation environment. We incorporate RaIR into free play and use it to complement the model’s epistemic uncertainty as an intrinsic reward. Doing so, we witness the autonomous construction of towers and other regular structures during free play, which leads to a substantial improvement in zero-shot downstream task performance on assembly tasks.

Regularity as Intrinsic Reward for Free Play

Abstract

We propose regularity as a novel reward signal for intrinsically-motivated reinforcement learning. Taking inspiration from child development, we postulate that striving for structure and order helps guide exploration towards a subspace of tasks that are not favored by naive uncertainty-based intrinsic rewards. Our generalized formulation of Regularity as Intrinsic Reward (RaIR) allows us to operationalize it within model-based reinforcement learning. In a synthetic environment, we showcase the plethora of structured patterns that can emerge from pursuing this regularity objective. We also demonstrate the strength of our method in a multiobject robotic manipulation environment. We incorporate RaIR into free play and use it to complement the model's epistemic uncertainty as an intrinsic reward. Doing so, we witness the autonomous construction of towers and other regular structures during free play, which leads to a substantial improvement in zero-shot downstream task performance on assembly tasks. Code and videos are available at https://sites.google.com/view/rair-project .

1 Introduction

Regularity, and symmetry as a specific form of regularity, are ubiquitous in nature as well as in our manufactured world. The ability to detect regularity helps to identify essential structures, minimizing redundancy and allowing for efficient interaction with the world [1]. Not only do we encounter symmetries in arts, design, and architecture, but our preference also showcases itself in play behavior. Adults and children have both been observed to prefer symmetry in visual perception, where symmetric patterns are more easily detected, memorized, and copied [2, 3]. Several works in developmental psychology show that regular patterns and symmetries are actively sought out during free play in children as well as adults [4-6].

Considering this in the context of a child's developmental cycle is intriguing. Studies show that children at the age of 2 exhibit a shift in their exploratory behavior. They progress from engaging in random actions on objects and unstable arrangements to purposefully engaging in functional activities and intentionally constructing stable configurations [7, 8, 5]. Bailey [4] reports that by 5 years of age, children build more structured arrangements out of blocks that exhibit alignment, balance, and examples of symmetries [5].

Despite the dominance of regularity in our perceptual systems and our preference for balance and stability during play, these principles are not yet well investigated within intrinsically-motivated reinforcement learning (RL). One prominent intrinsic reward definition is novelty, i.e. the agent is incentivized to visit areas of the state space with high expected information gain [9-12]. However, one fundamental problem with plain novelty-seeking objectives is that the search space is often unconstrained and too large. As an agent only has limited resources to allocate during play time, injecting appropriate inductive biases is crucial for sample efficiency, good coverage during exploration, and emergence of diverse behaviors. As proposed by Sancaktar et al. [12], using structured world models to inject a relational bias into exploration, yields more object and interaction-related novelty signals. However, which types of information to prioritize are not explicitly encoded in any of these methods. The direction of exploration is often determined by the inherent biases in the practical methods deployed. With imperfect world models that have a limited learning capacity and finite-horizon planning, novelty-seeking methods are observed to prefer 'chaotic' dynamics, where small perturbations lead to diverging trajectories, such as throwing, flipping, and poking objects. This in turn means that behaviors focusing on alignment, balance, and stability are overlooked. Not only are these behaviors relevant, as shown in developmental psychology, they also enable expanding and diversifying the types of behavior uncovered during exploration. As the behaviors observed during exploration are highly relevant for being able to solve downstream tasks, a chaos-favoring exploration will make it hard to solve assembly tasks, such as stacking. Indeed, successfully solving assembly tasks with more than 2 objects has been a challenge for intrinsically-motivated reinforcement learning.

We pose the question: how can we define an intrinsic reward signal such that RL agents prefer structured and regular patterns? We propose RaIR: R egularity a s I ntrinsic R eward, which aims to achieve highly ordered states. Mathematically, we operationalize this idea using entropy minimization of a suitable state description. Entropy and symmetries have been linked before [13, 14], however, we follow a general notion of regularity, i.e. where patterns reoccur and thus their description exhibits high redundancy / low entropy. In this sense, symmetries are a consequence of being ordered [15]. Regularity also means that the description is compressible, which is an alternative formulation. As argued by Schmidhuber [16], aiming for compression-progress is a formidable curiosity signal, however, it is currently unclear how to efficiently predict and optimize for it.

After studying the design choices in the mathematical formulation of regularity and the relation to symmetry operations, we set out to evaluate our regularity measure in the context of model-based reinforcement learning/planning, as it allows for highly sample-efficient exploration and solving complex tasks zero-shot, as shown in Sancaktar et al. [12]. To get a clear understanding of RaIR, we first investigate the generated structures when directly planning to optimize it using the ground truth system dynamics. A plethora of patterns emerge that are highly regular , as illustrated in Fig. 1.

Our ultimate goal is, however, to inject the proposed regularity objective into a free-play phase, where a robot can explore its capabilities in a task-free setting. During this free play, the dynamics model is learned on-the-go. We build on CEE-US [12], a free-play method that uses an ensemble of graph neural networks as a structured world model and the model's epistemic uncertainty as the only intrinsic reward. The epistemic uncertainty is estimated by the ensemble disagreement and acts as an effective novelty-seeking signal. We obtain structure-seeking free play by combining the conventional novelty-seeking objective with RaIR.

Our goal is to operationalize regularity, which is a well-established concept in developmental psychology, within intrinsically motivated RL. Furthermore, we showcase that biasing informationsearch towards regularity with RaIR indeed leads to the construction of diverse regular structures during play and significantly improves zero-shot performance in downstream tasks that also favor regularity, most notably assembly tasks. Besides conceptual work on compression [16, 17], to our knowledge, we are the first to investigate regularity as an intrinsic reward signal, bridging the gap between the diversity of behaviors observed in children's free play and what we can achieve with artificial agents.

Ensemble of

Update

agent's model

Generate

data

Maximize Regularity + Epistemic Uncertainty as intrinsic reward

RaIR

= - Entropy =

Plans in imagination for horizon

H

1

2

3

(a) RaIR + CEE-US: Regularity with Epistemic Uncertainty of Structured World Models

Relation Matrix

. . .

. . .

. . .

. . .

. . .

(b) RaIR computation

2 Method

First, we introduce our intrinsic reward definition for regularity. Then, we present its practical implementation and explain how we combine this regularity objective into model learning within free play.

2.1 Preliminaries

In this work, we consider environments that can be described by a fully observable Markov Decision Process (MDP), given by the tuple ( S , A , f a ss ′ , r a ss ′ ) , with the state-space S ∈ R n s , the action-space A ∈ R n a , the transition kernel f : S × A - →S , and the reward function r . Importantly, we consider the state-space to be factorized into the different entities, e.g. S = ( S obj ) N ×S robot for the state space of a robotic agent and N objects. We use model-based reinforcement learning, where data from interactions with the environment is used to learn a model ˜ f of the MDP dynamics [18]. Using this model, we consider finite-horizon ( H ) optimization/planning for undiscounted cumulative reward:

where s t + h are imagined states visited by rolling out the actions using ˜ f , which is assumed to be deterministic. The optimization of Eq. 1 is done with the improved Cross-Entropy Method (iCEM) [19] in a model predictive control (MPC) loop, i.e. re-planning after every step in the environment. Although this is not solving the full reinforcement learning problem (infinite horizon and stochastic environments), it is very powerful in optimizing for tasks on-the-fly and is thus suitable for optimizing changing exploration targets and solving downstream tasks zero-shot.

2.2 Regularity as Intrinsic Reward

Quite generally, regularity refers to the situation in which certain patterns reoccur. Thus, we formalize regularity as the redundancy in the description of the situation, to measure the degree of sub-structure recurrence. A decisive question is: which description should we use? Naturally, there is certain freedom in this choice, as there are many different coordinate frames. For instance, we could consider the list of absolute object positions or rather a relative representation of the scene.

To formalize, we define a mapping Φ : S → {X} + from state to a multiset {X} + of symbols (e.g. coordinates). A multiset is a set where elements can occur multiple times, e.g. { a, a, b } + . This multiset can equivalently be described by a tuple ( X,m ) , where X is the set of the unique elements, and m : X → Z + is a function assigning the multiplicity, i.e. the number of occurrences m ( x ) for

the elements x ∈ X . For the previous example, we get ( { a, b } , { a : 2 , b : 1 } ) . Given the multiset ( X,m ) ∈ {X} + , we define the discrete empirical distribution, also referred to as a histogram, by the relative frequency of occurrence p ( x ) = m ( x ) / ∑ x ′ ∈ X m ( x ′ ) for x ∈ X .

We define the regularity reward metric using (negative) Shannon entropy [20] of this distribution as:

We will now discuss concrete cases for the mapping Φ , i.e. how to describe a particular state.

Direct RaIR. In the simplest case, we describe the state s directly by the properties of each of the entities. For that, we define the function ϕ : S obj → {X} + , that maps each entity to a set of symbols and obtain Φ( s ) = ⋓ N i =1 ϕ ( s obj ,i ) as a union of all symbols. The symbols can be, for instance, discretized coordinates, colors, or other properties of the entities.

Let us consider the example where ϕ is extracting the object's Cartesian x and y coordinates in a rounded manner as ϕ ( s ) = {⌊ s x ⌉ , ⌊ s y ⌉} + , as shown in Fig. 3. The most irregular configuration would be when no two objects share the same rounded value in x and y . The object configuration becomes more and more regular the more objects share the same ⌊ x ⌉ and ⌊ y ⌉ coordinates. The most regular configuration is if all objects are in the same place. Note that this choice favors an axis-aligned configuration, and it is not invariant under global rotations.

0

6

16

13

10

3

0

occurrence

6

10 16

Relational RaIR of order k . Our framework for regularity quantification can easily be extended to a relational perspective, where we don't compute the entropy over aspects of individual entity properties, but instead on their pairwise or higher-order relations . This means that for a k -order regularity measure, we are interested in tuples of k entities. Thus, the mapping function ϕ no longer takes single entities as input, but instead operates on k -tuples:

ϕ is a function that describes some relations between the k input entities by a set of symbols.

For k -order regularity, the multiset Φ , over which we compute the entropy, is now given by

merged from all k -permutations of the N entities, denoted as P( { 1 , . . . , N } , k ) . In the case of a permutation invariant ϕ , Eq. 4 regards only the combinations C( { 1 , . . . , N } , k ) . Note that direct RaIR is equivalent to the relational RaIR of order 1 . Given the mapping Φ , the RaIR measure is computed as before with Eq. 2.

Depending on the order k and the function ϕ , we can select which regularities are going to be favored. Let us consider the example of a pairwise relational RaIR ( k = 2 ), where ϕ computes the relative positions: ϕ ( s i , s j ) = {⌊ s i -s j ⌉} , and rounding is performed elementwise. Whenever two entities have the same relative position to each other, the redundancy is detected. For k = 3 our regularity measure would be able to pick up sub-patterns composed of three objects, such as triangles and so forth.

As we are interested in physical interactions of the robot with objects and objects with objects, we choose RaIR of order k = 2 and explore various ϕ functions.

2.2.1 Properties of RaIR with Pairwise Relations and Practical Implementation

For simplicity, we are considering in the following that ϕ maps to a single symbol. Then for pairwise relationships ( k = 2 ), RaIR can be implemented using a relation matrix F ∈ X N × N . The entries F ij are given by ϕ ( s i , s j ) with s i , s j ∈ S obj. After constructing the relation matrix, we need the histogram of occurrences of unique values in this matrix to compute the entropy (Eq. 2). For continuous state spaces, the mapping function needs to implement a discretization step, which we implement by a binning of size b . For simplicity of notation, we reuse the rounding notation ⌊·⌉ for this discretization

(1) This is for one glide refl. operation. RaIR is increased for 2 glide refl. composition as it collapses onto transl.

step. This bin size b determines the precision of the measured regularity. In practice, we do not apply ϕ on the full entity state space, but on a subspace that contains e.g. the x -y (-z ) positions.

To understand the properties of our regularity measure for different ϕ , we present in Table 1 a categorization using the known symmetry operations in 2D and the following ϕ (applied to x -y positions): direct ϕ ( s i ) = ⌊ s i ⌉ (see previous section), relative position (difference vector) ϕ ( s i , s j ) = ⌊ s i -s j ⌉ , absolute value of the relative position 1 ϕ ( s i , s j ) = |⌊ s i -s j ⌉| , and Euclidean distance ϕ ( s i , s j ) = ⌊∥ s i -s j ∥⌉ . Figure 2b illustrates the RaIR computation using the absolute value of the relative position.

In Table 1, we first consider whether the measure is invariant under symmetry operations. That means if the value of RaIR stays unchanged when the entire configuration is transformed. We find that both Euclidean distance and relative position are invariant to all symmetry operations. The second and possibly more important question is whether a configuration with sub-structures of that symmetry has a higher regularity value than without, i.e. will patterns with these symmetries be favored. We find that Euclidean distance favors all symmetries, followed by absolute value of the relative position. A checkmark in this part of the table means that the more entities can be mapped to each other with the same transformation, the higher RaIR. Although the Euclidean distance seems favorable, we find that it mostly clumps entities and creates fewer alignments. To get a sense of the patterns scoring high in the regularity measure, Fig. 1 showcases situations that emerge when RaIR with absolute value of relative position is optimized (details below).

2.3 Regularity in Free Play

Our goal is to explicitly put the bias of regularity into free play via RaIR, as in Fig. 2a. What we want to achieve is not just that the agent creates regularity, but that it gathers valuable experience in creating regularity. This ideally leads to directing exploration towards patterns/arrangements that are novel.

We propose to use RaIR to augment plain novelty-seeking intrinsic rewards, in this work specifically ensemble disagreement [10]. We choose ensemble disagreement because 1) we need a reward definition that allows us to predict future novelty, such that we can use it inside model-based planning (this constraint makes methods relying on retrospective novelty such as Intrinsic Curiosity Module (ICM) [9] ineligible), and 2) we want to use the models learned during free play for zero-shot downstream task generalization via planning in a follow-up extrinsic phase. It has been shown in previous works that guiding exploration by the model's own epistemic uncertainty, approximated

via ensemble disagreement, leads to learning more robust world models compared to e.g. Random Network Distillation (RND) [21] (Suppl. J.1), resulting in improved zero-shot downstream task performance [12]. That is why we choose ensemble disagreement to compute expected future novelty.

We train an ensemble of world models { ( ˜ f θm ) M m =1 } , where M denotes the ensemble size. The model's epistemic uncertainty is approximated by the disagreement of the ensemble members' predictions. The disagreement reward is given by the trace of the covariance matrix [12]:

We incorporate our regularity objective into free play by using a linear combination of RaIR and ensemble disagreement. Overall, we have the intrinsic reward:

where λ controls the trade-off between regularity and pure epistemic uncertainty.

Model-based Planning with Structured World Models To optimize the reward function on-thefly, we use model-based planning using zero-order trajectory optimization, as introduced in Sec. 2.1. Concretely, we use CEE-US [12], which combines structured world models and epistemic uncertainty (Eq. 5) as intrinsic reward. The structured world models are ensembles of message-passing Graph Neural Networks (GNNs) [22], where each object corresponds to a node in the graph. The node attributes { s t,i ∈ S obj | i = 1 , . . . , N } are the object features such as position, orientation, and velocity at time step t . The state representation of the actuated agent s robot ∈ S robot similarly contains position and velocity information about the robot. We treat the robot as a global node in the graph [12]. We refer to the combination of RaIR with ensemble disagreement, medium-horizon planning (20-30 time steps), and structured world models as RaIR + CEE-US.

3 Experiments

We evaluate RaIR in the two environments shown in Fig. 1.

ShapeGridWorld is a grid environment, where each circle represents an entity/agent that is controlled separately in x -y directions. Entities are controlled one at a time. Starting from timestep t = 0 , the entity with i = 1 is actuated for T timesteps, where T is the entity persistency. Then, at t = T , actuation switches over to entity i = 2 and we keep iterating over the entities in this fashion. Each circle is treated as an entity for the regularity computation with a 2D-entity state space S obj with x -y positions.

Fetch Pick & Place Construction is an extension of the Fetch Pick & Place environment [23] to more cubes [24] and a large table [12]. An end-effector-controlled robot arm is used to manipulate blocks. The robot state S robot ∈ R 10 contains the end-effector position and velocity, and the gripper's state (open/close) and velocity. Each object's state S obj ∈ R 12 is given by its pose and velocities. For free play, we use 6 objects and consider several downstream tasks with varying object numbers.

3.1 Emerging Patterns in SHAPEGRIDWORLD and CONSTRUCTION with RaIR

To get a sense of what kinds of patterns emerge following our regularity objective with RaIR, we do planning using ground truth (GT) models, i.e. with access to the true simulator itself for planning. We perform these experiments to showcase that we can indeed get regular constellations with our proposed formulation. Since we can perform multi-horizon planning without any accumulating model errors using ground truth models, we can better investigate the global/local optima of our regularity reward. Note that as we are using a zero-order trajectory optimizer with a limited sample budget and finite-horizon planning, we don't necessarily converge to the global optima. We use ϕ ( s i , s j ) = { ( |⌊ s i,x -s j,x ⌉| , |⌊ s i,y -s j,y ⌉| ) } for RaIR in both environments. The emerging patterns are shown in Fig. 1.

In the 2D SHAPEGRIDWORLD environment, we indeed observe that regular patterns with translational, reflectional (axis-aligned), glide-reflectional (axis-aligned), and rotational symmetries emerge. Regularity is not restricted to spatial relations and can be applied to any set of symbols. To showcase this, we apply RaIR to colored SHAPEGRIDWORLD, where color is part of S obj. Generated patterns are not just regular in x -y but also in color, as shown in Fig. S7 and the 2 top right examples in Fig. 1.

For CONSTRUCTION, we also observe complex constellations with regularities, even stacks of all 6 objects. Since we are computing RaIR on the x -y positions, a stack of 6 is the global optimum. The optimization of RaIR for this case is shown in Fig. 4. Note that stacking itself is a very challenging task, and was so far only reliably achievable with reward shaping or tailored learning curricula [24]. The fact that these constellations appear naturally from our regularity objective, achievable with a planning horizon of 30 timesteps, is by itself remarkable.

0

100

200

300

Timestep

t

-

3

-

2

-

1

0

RaIR

Additional example patterns generated in CONSTRUCTION with RaIR on the x -y -z positions can be found in the Suppl. A. In that case, a horizontal line on the ground and a vertical line into air, i.e. a stack, are numerically equivalent with respect to RaIR. Choosing to operate on the x -y -subspace is injecting the direction of gravity and provides a bias towards vertical alignments. We also apply RaIR to a custom CONSTRUCTION environment with different shapes and masses (cubes, columns, balls and flat blocks) and once again observe regular arrangements, as in Fig. S8. More details in Suppl. F.

3.2 Free Play with RaIR in CONSTRUCTION

We perform free play in CONSTRUCTION, i.e. only optimize for intrinsic rewards, where we learn models on-the-go. During free play, we start with randomly initialized models and an empty replay buffer. Each iteration of free play consists of data collection with environment interactions (via online planning), and then model training on the collected data so far (offline).

In each iteration of free play, we collect 2000 samples (20 rollouts with 100 timesteps each) and add them to the replay buffer. During the online planning part for data collection, we only perform inference with the models and no training is performed. Afterwards, we train the model for a fixed number of epochs on the replay buffer. We then continue with data collection in the next free play iteration. More details can be found in Suppl. H. For this intrinsic phase, we combine our regularity objective with ensemble disagreement as per Eq. 6. The goal is to bias exploration and the search for information gain towards regular structures, corresponding to the optima that emerge with ground truth models, as shown in Fig. 1. We also show results for the baselines RND [21] and Disagreement (Dis) [10], using the same model-based planning backbone as CEE-US (see Suppl. J). The Dis baseline also uses ensemble disagreement as intrinsic reward, however unlike CEE-US, only plans for one step into the future.

In Figure 5, we analyze the quality of data generated during free play, in terms of observed interactions, for RaIR + CEE-US with the augmentation weight λ = 0 . 1 , a pure RaIR run with no information-gain component in the intrinsic reward ( λ = 0 ), CEE-US, as well as RND and Disagreement. For pure RaIR, we observe a decrease in the generated interactions. This has two reasons: 1) RaIR only aims to generate structure and the exploration problem is not solved, 2) once the controller finds a plan that leads to an optimum, even if it is local, there is no incentive to destroy

RaIR

RaIR+CEE-US

CEE-US

RND

Dis

relative time

0

50

100

150

200

250

300

Free Play Iteration

0%

25%

50%

(a) 1 object moves

0

50

100

150

200

250

300

Free Play Iteration

0%

25%

50%

(b) 2 or more objects move

0

50

100

150

200

250

300

Free Play Iteration

0%

15%

30%

(c) object(s) in air

0

50

100

150

200

250

300

Free Play Iteration

0%

15%

30%

(d) object(s) flipped

Iteration 251

Iteration 255

Iteration 259

Iteration 275

Iteration 284

Iteration 286

Iteration 296

it, unless a plan that results in better regularity can be found within the planning horizon. There is no discrimination between 'boring' and 'interesting' patterns with respect to the model's current capabilities. This in turn means that the robot creates e.g. a (spaced) line, which is a local optimum for RaIR, and then spends the rest of the episode, not touching any objects to keep the created alignment intact. With the injection of some disagreement in RaIR + CEE-US, we observe improved interaction metrics throughout free play in terms of 2 or more object interactions and objects being in the air (either being lifted by the robot or being stacked on top of another block). In practice, since the ensemble of models tends to hallucinate due to imperfect predictions, even for pure RaIR we observe dynamic pattern generations, as reflected in the interaction metrics (more details in Suppl. C). For the plain disagreement case with CEE-US, more flipping behavior, and less air time are observed during free play, since the agent favors chaos. We also observe that the baselines RND and Disagreement (planning horizon 1) produce less interaction-rich data during free play. Especially for Disagreement, this further shows that planning for future novelty is an essential component for free play.

Another reason why disagreement in RaIR + CEE-US is helpful is due to the step-wise landscape of RaIR as shown in Fig. 4. Here, combining RaIR with ensemble disagreement effectively helps smoothen this reward function, making it easier to find plans with improvements in regularity with imperfect world models.

In Fig. 7, we report the highest achieved RaIR value in the collected rollouts throughout free play. We observe that pure RaIR and RaIR+CEE-US indeed find more regular structures during play compared to the baselines. Some snapshots of regular structures observed during a RaIR + CEE-US free play run are illustrated in Fig. 6. Results for ϕ ( s i , s j ) = ⌊ s i -s j ⌉ can be found in Suppl. B.

0

100

200

300

Free Play Iteration

-2.50

-2.25

-2.00

-1.75

RaIR

3.3 Zero-shot Generalization to Assembly Downstream Tasks with RaIR in CONSTRUCTION

After the fully-intrinsic free-play phase, we evaluate zero-shot generalization performance on downstream tasks, where we perform model-based planning with the learned world models. Note that now instead of optimizing for intrinsic rewards, we are optimizing for extrinsic reward functions r task given by the environment (Suppl. I.4.1). In Fig. 8, we present the evolution of success rates of models checkpointed throughout free play on the following assembly tasks: singletower with 3 objects, 2 multitowers with 2 objects each, pyramid with 5 and 6 objects.

The combination RaIR + CEE-US yields significant improvements in the success rates of assembly tasks, as shown in Fig. 8 and Table 2. RaIR alone, outperforms CEE-US in assembly tasks. As we are biasing exploration towards regularity, we see a decrease in more chaotic interactions during play time, which is correlated with a decrease in performance for the more chaotic throwing and flipping tasks. For the generic Pick & Place task, we observe comparable performance between RaIR+CEE-US and CEE-US. The decrease in performance for RaIR in non-assembly tasks shows the importance of an information-gain objective. The baselines RND and Disagreement exhibit very poor performance in the assembly tasks. In the other tasks, RND and CEE-US are comparable (bold numbers show statistical indistinguishability from best with p > 0 . 05 ) This showcases that guiding free play by the model's epistemic uncertainty as in the case with ensemble disagreement, helps learn robust and capable world models. The decrease in the zero-shot performance for Disagreement further proves the importance of planning for future novelty during free play. We also run free play in the custom CONSTRUCTION environment for RaIR + CEE-US and CEE-US and observe improved zero-shot downstream performance in assembly tasks with RaIR + CEE-US (see Suppl. F.2).

success rate

Singletower 3

0

100

200

300

Free Play Iteration

0%

50%

100%

Multitower 2+2

0

100

200

300

Free Play Iteration

0%

50%

100%

Pyramid 5

0

100

200

300

Free Play Iteration

0%

30%

60%

Pyramid 6

0

100

200

300

Free Play Iteration

0%

12%

25%

RaIR

RaIR+CEE-US

CEE-US

RND

Dis

3.4 Re-creating existing structures with RaIR

We test whether we can re-create existing arrangements in the environment with RaIR. If there are regularities / sub-structures already present in the environment, then completing or re-creating these patterns naturally becomes an optimum for RaIR, as repeating this pattern introduces redundancy. We initialize pyramids, single- and multitowers out of the robot's manipulability range in CONSTRUCTION. We then plan using iCEM to maximize RaIR with GT models. Doing so, the agent manages to re-create the

(a)

t

= 0

(b)

t

= 200

existing structures in the environment with the blocks it has within reach. In Fig. 9, this is showcased for a pyramid with 3 objects, where in 15 rollouts a pyramid is recreated in 73% of the cases. Without the need to define any explicit reward functions, we can simply use our regularity objective to mimic existing ordered constellations. More details can be found in Suppl. D.

4 RaIR in Locomotion Environments

The only requirement to incorporate RaIR in a given environment is an entity-based view of the world. We can easily apply this principle to locomotion environments, where we treat each joint as an entity. In Figure 10, we showcase the generated poses in the DeepMind Control Suite environments Quadruped and Walker, when we optimize for regularity using ground truth models. Here, regularity is computed over the x -y -z positions of the knees and toes of the quadruped. For walker, we take the positions of the feet, legs and the torso. These poses also heavily overlap with the goal poses proposed in the RoboYoga benchmark [25] (see Fig. S11), which further supports our hypothesis that regularities and our preference for them are ubiquitious. We also apply free play with learned models in the Quadruped environment. More details and results can be found in Suppl. G.

Walker

5 Related Work

Intrinsic motivation in RL uses minimizing novelty/surprise to dissolve cognitive disequilibria as a prominent intrinsic reward signal definition [9, 26-31]. As featured in this work, using the disagreement of an ensemble of world models as an estimate of expected information gain is a widely-used metric as it allows planning into the future [10-12]. Other prominent intrinsic rewards deployed in RL include learning progress [26, 30, 32], empowerment [33, 34] and maximizing for state space coverage with count-based methods [35, 36] and RND [21]. Another sub-category would be goal-conditioned unsupervised exploration methods combined with e.g. ensemble disagreement [25, 37] or asymmetric self-play [38]. In competence-based intrinsic motivation, unsupervised skill discovery methods aim to learn policies conditioned on a latent skill variable [39, 40]. Berseth et al. [41] propose surprise minimization as intrinsic reward to seek familiar states in unstable environments with an active source of entropy. Note that this differs from our work, as our notion of regularity is decoupled from surprise: RaIR aims to get to a state that is regular in itself.

Compression and more specifically compression progress have been postulated as driving forces in human curiosity by Schmidhuber [16]. However, the focus has been on temporal compression, where it is argued that short and simple explanations of the past make long-horizon planning easier. In our work, we don't focus on compression in the temporal dimension, i.e. sequences of states. Instead, we perform compression as entropy minimization (in the relational case, equivalent to lossy compression) at a given timestep t , where we are interested in the relational redundancies in the current scene. More details on connections to compression can be found in Appendix. L.

Assembly Tasks in RL with 3+ objects pose an open challenge, where most methods achieve stacking via tailored learning curricula with more than 20 million environment steps [24, 42], expert demonstrations [43], also together with high-level actions [44]. Hu et al. [37] manage to solve 3-object stacking in an unsupervised setting with goal-conditioned RL (GCRL), using a very similar robotic setup to ours, but only with 30% success rate. More discussion on GCRL in Suppl. K.

6 Discussion

Although the search for regularity and symmetry has been studied extensively in developmental psychology, these concepts haven't been featured within reinforcement learning yet. In this work, we propose a mathematical formulation of regularity as an intrinsic reward signal and operationalize it within model-based RL. We show that with our formulation of regularity, we indeed manage to create regular and symmetric patterns in a 2D grid environment as well as in a challenging compositional object manipulation environment. We also provide insights into the different components of RaIR and deepen the understanding of the types of regularities emerging from using different mappings ϕ . In the second part of the work, we incorporate RaIR within free play. Here, our goal is to bias information-search during exploration towards regularity. We provide a proof-of-concept that augmenting epistemic uncertainty-based intrinsic rewards with RaIR helps exploration for symmetric and ordered arrangements. Finally, we also show that our regularity objective can simply be used to imitate existing regularities in the environment.

Limitations and future work: As we use finite-horizon planning, we don't necessarily converge to global optima. This can both be seen as a limitation and a feature, as it naturally allows us to obtain different levels of regularity in the generated patterns. Currently, we are restricted to fully-observable MDPs. We embrace object-centric representations as a suitable inductive bias in RL, where the observations per object (consisting of poses and velocities) are naturally disentangled (more discussion in Appendix. M). We also assume that this state space is interpretable such that we take, for instance, only the positions and color. The representational space, in which the RaIR measure is computed, is specified by the designer. Exciting future work would be to learn a representation under which the human relevant structures in the real-world (e.g. towers, bridges) are indeed regular.

Acknowledgments and Disclosure of Funding

The authors thank Sebastian Blaes, Anselm Paulus, Marco Bagatella, Núria Armengol Urpí and Maximilian Seitzer for helpful discussions and for their help reviewing the manuscript. The authors thank the International Max Planck Research School for Intelligent Systems (IMPRS-IS) for supporting Cansu Sancaktar. 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).

References

Yanxi Liu, Hagit Hel-Or, Craig S Kaplan, Luc Van Gool, et al. Computational symmetry in computer vision and computer graphics. Foundations and Trends® in Computer Graphics and Vision , 5(1-2):1-195, 2010.

Fred Attneave. Symmetry, information, and memory for patterns. The American journal of psychology , 68(2):209-222, 1955.

Marc H Bornstein, Kay Ferdinandsen, and Charles G Gross. Perception of symmetry in infancy. Developmental psychology , 17(1):82, 1981.

Marjory W Bailey. A scale of block constructions for young children. Child Development , 4(2): 121-139, 1933.

William A Zingrone. The construction of symmetry in children and adults. The Journal of genetic psychology , 175(2):91-104, 2014.

Claire Golomb. The development of compositional strategies in children's drawings. Visual Arts Research , pages 42-52, 1987.

Eleanor J Gibson. Exploratory behavior in the development of perceiving, acting, and the acquiring of knowledge. Annual review of psychology , 39(1):1-42, 1988.

Jonas Langer. The origins of logic: One to two years . Academic Press, 1986. ISBN 9780124365551.

Deepak Pathak, Pulkit Agrawal, Alexei A. Efros, and Trevor Darrell. Curiosity-driven exploration by self-supervised prediction. In International Conference on Machine Learning (ICML) , 2017. URL https://proceedings.mlr.press/v70/pathak17a.html .

Deepak Pathak, Dhiraj Gandhi, and Abhinav Gupta. Self-supervised exploration via disagreement. In International Conference on Machine Learning (ICML) , 2019. URL https: //proceedings.mlr.press/v97/pathak19a.html .

Ramanan Sekar, Oleh Rybkin, Kostas Daniilidis, Pieter Abbeel, Danijar Hafner, and Deepak Pathak. Planning to explore via self-supervised world models. In International Conference on Machine Learning (ICML) , 2020. URL https://proceedings.mlr.press/v119/ sekar20a.html .

Cansu Sancaktar, Sebastian Blaes, and Georg Martius. Curious exploration via structured world models yields zero-shot object manipulation. In Advances in Neural Information Processing Systems (NeurIPS) , 2022.

Shu-Kun Lin. Correlation of entropy with similarity and symmetry. Journal of Chemical Information and Computer Sciences , 36(3):367-376, 1996.

Daniel Lazarev. Information measures for entropy and symmetry. arXiv preprint arXiv:2211.14857 , 2022.

Edward Bormashenko. Entropy, information, and symmetry; ordered is symmetrical, ii: system of spins in the magnetic field. Entropy , 22(2):235, 2020.

Jürgen Schmidhuber. Driven by compression progress: A simple principle explains essential aspects of subjective beauty, novelty, surprise, interestingness, attention, curiosity, creativity, art, science, music, jokes. In Anticipatory Behavior in Adaptive Learning Systems: From Psychological Theories to Artificial Cognitive Systems 4 , pages 48-76. Springer, 2009. Longer preprint: https://arxiv.org/abs/0812.4360 .

Jürgen Schmidhuber. Powerplay: Training an increasingly general problem solver by continually searching for the simplest still unsolvable problem. Frontiers in Psychology , 4:313, 2013.

Kurtland Chua, Roberto Calandra, Rowan McAllister, and Sergey Levine. Deep reinforcement learning in a handful of trials using probabilistic dynamics models. In Advances in Neural Information Processing Systems (NeurIPS) , 2018. URL https://proceedings.neurips.cc/ paper_files/paper/2018/file/3de568f8597b94bda53149c7d7f5958c-Paper.pdf .

Cristina Pinneri, Shambhuraj Sawant, Sebastian Blaes, Jan Achterhold, Joerg Stueckler, Michal Rolinek, and Georg Martius. Sample-efficient cross-entropy method for real-time planning. In Conference on Robot Learning (CoRL) , 2020. URL https://proceedings.mlr.press/ v155/pinneri21a.html .

Claude E Shannon. A mathematical theory of communication. The Bell system technical journal , 27(3):379-423, 1948.

Yuri Burda, Harrison Edwards, Amos Storkey, and Oleg Klimov. Exploration by random network distillation. In International Conference on Learning Representations (ICLR) , 2019. URL https://openreview.net/forum?id=H1lJJnR5Ym .

Peter W. Battaglia, Jessica B. Hamrick, Victor Bapst, Alvaro Sanchez-Gonzalez, Vinicius Zambaldi, Mateusz Malinowski, Andrea Tacchetti, David Raposo, Adam Santoro, Ryan Faulkner, Caglar Gulcehre, Francis Song, Andrew Ballard, Justin Gilmer, George Dahl, Ashish Vaswani, Kelsey Allen, Charles Nash, Victoria Langston, Chris Dyer, Nicolas Heess, Daan Wierstra, Pushmeet Kohli, Matt Botvinick, Oriol Vinyals, Yujia Li, and Razvan Pascanu. Relational inductive biases, deep learning, and graph networks. arXiv:1806.01261 , 2018. URL https://arxiv.org/abs/1806.01261 .

Matthias Plappert, Marcin Andrychowicz, Alex Ray, Bob McGrew, Bowen Baker, Glenn Powell, J. Schneider, Joshua Tobin, Maciek Chociej, P. Welinder, V. Kumar, and W. Zaremba. Multi-goal reinforcement learning: Challenging robotics environments and request for research. arXiv preprint arXiv:1802.09464 , 2018. URL https://arxiv.org/abs/1802.09464 .

Richard Li, Allan Jabri, Trevor Darrell, and Pulkit Agrawal. Towards practical multi-object manipulation using relational reinforcement learning. In 2020 IEEE International Conference on Robotics and Automation (ICRA) , pages 4051-4058. IEEE, 2020. URL https://arxiv. org/abs/1912.11032 .

Russell Mendonca, Oleh Rybkin, Kostas Daniilidis, Danijar Hafner, and Deepak Pathak. Discovering and achieving goals via world models. In Advances in Neural Information Processing Systems (NeurIPS) , 2021. URL https://openreview.net/forum?id=6vWuYzkp8d .

Jürgen Schmidhuber. A possibility for implementing curiosity and boredom in model-building neural controllers. In Proceedings of the International Conference on Simulation of Adaptive Behavior: From Animals to Animats , 1991. URL https://dl.acm.org/doi/10.5555/ 116517.116542 .

Kuno Kim, Megumi Sano, Julian De Freitas, Nick Haber, and Daniel Yamins. Active world model learning with progress curiosity. In International Conference on Machine Learning (ICML) , 2020. URL https://arxiv.org/abs/2007.07853 .

Oliver Groth, Markus Wulfmeier, Giulia Vezzani, Vibhavari Dasagi, Tim Hertweck, Roland Hafner, Nicolas Heess, and Martin Riedmiller. Is curiosity all you need? on the utility of emergent behaviours from curious exploration. arXiv preprint arXiv:2109.08603 , 2021. URL https://arxiv.org/abs/2109.08603 .

J. Storck, S. Hochreiter, and J. Schmidhuber. Reinforcement driven information acquisition in non-deterministic environments. In Proceedings of the International Conference on Artificial Neural Networks , pages 159-164, Paris, 1995. EC2 & Cie. URL https://people.idsia. ch/~juergen/icann95new.pdf .

Sebastian Blaes, Marin Vlastelica, Jia-Jie Zhu, and Georg Martius. Control What You Can: Intrinsically motivated task-planning agent. In Advances in Neural Information Processing Systems (NeurIPS) , 2019. URL https://proceedings.neurips.cc/paper/2019/hash/ b6f97e6f0fd175613910d613d574d0cb-Abstract.html .

Giuseppe Paolo, Alexandre Coninx, Stephane Doncieux, and Alban Laflaquière. Sparse reward exploration via novelty search and emitters. In Proceedings of the Genetic and Evolutionary Computation Conference , page 154-162, 2021. URL https://doi.org/10.1145/3449639. 3459314 .

Cédric Colas, Pierre Fournier, Mohamed Chetouani, Olivier Sigaud, and Pierre-Yves Oudeyer. CURIOUS: Intrinsically motivated modular multi-goal reinforcement learning. In International Conference on Machine Learning (ICML) , 2019. URL https://proceedings.mlr.press/ v97/colas19a.html .

A.S. Klyubin, D. Polani, and C.L. Nehaniv. Empowerment: a universal agent-centric measure of control. In IEEE Congress on Evolutionary Computation , volume 1, pages 128-135 Vol.1, 2005. URL https://ieeexplore.ieee.org/document/1554676 .

Shakir Mohamed and Danilo Jimenez Rezende. Variational information maximisation for intrinsically motivated reinforcement learning. In Advances in Neural Information Processing Systems (NeurIPS) , 2015. URL https://arxiv.org/abs/1509.08731 .

Marc Bellemare, Sriram Srinivasan, Georg Ostrovski, Tom Schaul, David Saxton, and Remi Munos. Unifying count-based exploration and intrinsic motivation. Advances in Neural Information Processing Systems (NeurIPS) , 2016.

Haoran Tang, Rein Houthooft, Davis Foote, Adam Stooke, OpenAI Xi Chen, Yan Duan, John Schulman, Filip DeTurck, and Pieter Abbeel. # exploration: A study of count-based exploration for deep reinforcement learning. Advances in Neural Information Processing Systems (NeurIPS) , 2017.

Edward S. Hu, Richard Chang, Oleh Rybkin, and Dinesh Jayaraman. Planning goals for exploration. In International Conference on Learning Representations (ICLR) , 2023. URL https://openreview.net/forum?id=6qeBuZSo7Pr .

OpenAI, Matthias Plappert, Raul Sampedro, Tao Xu, Ilge Akkaya, Vineet Kosaraju, Peter Welinder, Ruben D'Sa, Arthur Petron, Henrique Ponde de Oliveira Pinto, Alex Paino, Hyeonwoo Noh, Lilian Weng, Qiming Yuan, Casey Chu, and Wojciech Zaremba. Asymmetric selfplay for automatic goal discovery in robotic manipulation. arXiv:2101.04882 , 2021. URL https://arxiv.org/abs/2101.04882 .

Benjamin Eysenbach, Abhishek Gupta, Julian Ibarz, and Sergey Levine. Diversity is all you need: Learning skills without a reward function. In International Conference on Learning Representations (ICLR) , 2019. URL https://openreview.net/forum?id=SJx63jRqFm .

Archit Sharma, Shixiang Gu, Sergey Levine, Vikash Kumar, and Karol Hausman. Dynamicsaware unsupervised discovery of skills. In International Conference on Learning Representations (ICLR) , 2020. URL https://openreview.net/forum?id=HJgLZR4KvH .

Glen Berseth, Daniel Geng, Coline Manon Devin, Nicholas Rhinehart, Chelsea Finn, Dinesh Jayaraman, and Sergey Levine. {SM}irl: Surprise minimizing reinforcement learning in unstable environments. In International Conference on Learning Representations (ICLR) , 2021. URL https://openreview.net/forum?id=cPZOyoDloxl .

John B. Lanier, Stephen McAleer, and Pierre Baldi. Curiosity-driven multi-criteria hindsight experience replay, 2019.

Ashvin Nair, Bob McGrew, Marcin Andrychowicz, Wojciech Zaremba, and Pieter Abbeel. Overcoming exploration in reinforcement learning with demonstrations. In 2018 IEEE international conference on robotics and automation (ICRA) , pages 6292-6299. IEEE, 2018.

Ondrej Biza, Thomas Kipf, David Klee, Robert Platt, Jan-Willem van de Meent, and Lawson LS Wong. Factored world models for zero-shot generalization in robotic manipulation. arXiv preprint arXiv:2202.05333 , 2022. URL https://arxiv.org/abs/2202.05333 .

Yuval Tassa, Yotam Doron, Alistair Muldal, Tom Erez, Yazhe Li, Diego de Las Casas, David Budden, Abbas Abdolmaleki, Josh Merel, Andrew Lefrancq, et al. Deepmind control suite. arXiv preprint arXiv:1801.00690 , 2018.

Francesco Locatello, Dirk Weissenborn, Thomas Unterthiner, Aravindh Mahendran, Georg Heigold, Jakob Uszkoreit, Alexey Dosovitskiy, and Thomas Kipf. Object-centric learning with slot attention. Advances in Neural Information Processing Systems (NeurIPS) , 2020.

Pedro A. Tsividis, Joao Loula, Jake Burga, Nathan Foss, Andres Campero, Thomas Pouncy, Samuel J. Gershman, and Joshua B. Tenenbaum. Human-level reinforcement learning through theory-based modeling, exploration, and planning. arXiv:2107.12544 , 2021. URL https: //arxiv.org/abs/2107.12544 .

Peter Battaglia, Razvan Pascanu, Matthew Lai, Danilo Jimenez Rezende, et al. Interaction networks for learning about objects, relations and physics. Advances in Neural Information Processing Systems (NeurIPS) , 2016.

Elizabeth S Spelke. Principles of object perception. Cognitive science , 14(1):29-56, 1990.

Elizabeth S Spelke, Karen Breinlinger, Kristen Jacobson, and Ann Phillips. Gestalt relations and object perception: A developmental study. Perception , 22(12):1483-1501, 1993.

Benjamin Peters and Nikolaus Kriegeskorte. Capturing the objects of vision with neural networks. Nature human behaviour , 5(9):1127-1144, 2021.

Supplementary Material for Regularity as Intrinsic Reward for Free Play

Code and videos are available at https://sites.google.com/view/rair-project .

A Experiment Results with Ground Truth Models

A.1 Experiment Results for RaIR in CONSTRUCTION with x -y -z

As discussed in Sec. 3.1, in our experiments we compute RaIR on the x -y subspace of the object positions in CONSTRUCTION to inject a bias towards vertical alignments. Examples of patterns generated when optimizing for RaIR using

are showcased in Fig. S1. When we also include the z -positions of the objects in the RaIR computation, patterns and constellations on the ground are preferred. In this case, there is no difference between a horizontal line on the ground vs. a vertical line, i.e. a stack. Since creating a stack, however, is a more sparse solution, in practice the zero-order trajectory optimizer converges already to regular structures on the ground and vertical constellations don't emerge. Starting from a randomly initialized scene with all objects on the ground, the regularity metric for x -y -z only starts increasing when multiple objects are in the stack, which would require a very long planning horizon to find this solution.

B Experiment Results for Relational Case without Absolute Value

We present the interaction metrics observed during free play with RaIR + CEE-US in the case of relational ϕ with ϕ ( s i , s j ) = {⌊ s i -s j ⌉} = { ( ⌊ s i,x -s j,x ⌉ , ⌊ s i,y -s j,y ⌉ ) } . We find the interaction metrics to be comparable to the absolute relational case presented in the main paper with ϕ ( s i , s j ) = { ( |⌊ s i,x -s j,x ⌉| , |⌊ s i,y -s j,y ⌉| ) } . In Fig. S2, we also include the results for RaIR+CEE-US with λ = 1 . In the case of the increased weighting on the ensemble disagreement term, the free-play behavior indeed collapses back onto CEE-US. This means we have more flipping behavior and less air time. In the case of RaIR + CEE-US with smaller λ = 0 . 1 , we seek regular states, which include vertical alignments, such that the air time doesn't go down. Note that in this case, there is still an incentive to 'destroy' and lean towards chaos due to the ensemble disagreement reward term, such that the constellations showcased in Fig. 6 (snapshots for the absolute relational ϕ ) and Fig. S3 are not necessarily preserved.

We also evaluate the success rates for zero-shot downstream task generalization using the models trained in free-play runs with relational (R) ϕ and present them in Table S1. We find the performance in this case to be comparable to the absolute relational (AR) ϕ case. Note that in comparison to previous work [12], we use controller mode 'best', not 'sum' and a shorter planning horizon of 20 timesteps instead of 30. The environment is also intialized with 6 objects instead of 4. That's why the reported success rates differ from previous work. We also include the success rates for CEE-US with mode 'sum', and see that using mode 'best' for free play also results in improved performance in assembly tasks.

RaIR

RaIR+CEE-US

(

λ

= 0

.

1)

RaIR+CEE-US

(

λ

= 1)

CEE-US

relative time

0

50

100

150

200

250

300

Transitions × 2000

0%

25%

50%

(a) 1 object moves

0

50

100

150

200

250

300

Transitions × 2000

0%

25%

50%

(b) 2 or more objects move

0

100

200

300

Transitions × 2000

0%

20%

40%

(c) object(s) in air

0

50

100

150

200

250

300

Transitions × 2000

0%

15%

30%

(d) object(s) flipped

Iteration 245

Iteration 246

Iteration 247

Iteration 248

Iteration 270

Iteration 275

Iteration 295

C Experiment Results for Free Play with pure RaIR

In this section, we further discuss the zero-shot downstream task generalization performance for free play with pure RaIR and the role of the information-gain term in our intrinsic reward combination used in free play, as specified in Eq. 6. As discussed in Sec. 3.2, adding ensemble disagreement to our regularity objective leads to 1) more interaction-rich free play and 2) more robust world models which yield higher success rates for zero-shot downstream task generalization. For both the absolute relational case presented in Fig. 5 and the relational case in Fig. S2, RaIR with no disagreement term yields less interactions in terms of object(s) being moved, being in air and being flipped. This is because the exploration problem is not solved by RaIR alone. When we use ensemble disagreement as an intrinsic reward, the discovery of different types of interactions is accelerated. When one of the models in the ensemble learns a new type of dynamics, such as an object moving, the ensemble disagreement goes up, incentivizing the agent to repeat this behavior until it is learned by all models such that disagreement goes down. In the case of RaIR, this only happens implicitly: with some

Iteration 270

Iteration 273

Iteration 276

Iteration 295

Iteration 296

Iteration 298

Iteration 299

models in the ensemble learning a certain type of dynamics in the environment, during planning, the models can hallucinate objects being aligned and creating a regular pattern with high RaIR such that these actions are executed by the controller. These false attempts also help exploration.

As the models produce better predictions, especially after free play iteration 200, we observe that more stable patterns are generated with RaIR compared to RaIR + CEE-US (Fig. 7) and the amount of time objects are moving starts decreasing Fig. 5. This is because in this case when the models get better and hallucinate less, there is no reason to leave local optima such as a spaced line unless a pattern that yields a higher regularity value can be found within the planning horizon.

The challenge of exploration with pure RaIR is also reflected in the interaction time for object(s) in air. Starting to create regular patterns such as stacks takes longer, as exploring to lift objects happens later without the disagreement reward. This is also connected to the step-wise landscape of RaIR as discussed in Sec. 3.2 such that explicit exploration via ensemble disagreement is beneficial.

As showcased in Fig. S4, we still observe the stable generation of patterns such as spaced lines later on in training as well, as these are local optima of RaIR. However, we start to see more stacks generated in the later stages of free play. In Fig. 7, the highest RaIR value achieved for the different variants are showcased throughout training. Pure RaIR, achieves slightly higher regularity then RaIR + CEE-US. This is also because pure RaIR, tends to generate more regular patterns that feature all objects, i.e. all objects are in-line or build a square. With RaIR + CEE-US, as some chaotic behavior is injected to free play via ensemble disagreement, more local regularities such as a stack of 2, with the rest of the objects in disorder, are likely to emerge.

Through injecting ensemble disagreement into free play, the robustness of the learned world models is also increased as they are guided by their own epistemic uncertainty [12]. During free play, data is actively collected from regions where the models are uncertain, acting as their own adversary. This in turn makes the models more robust for deployment in model-based planning in the followup extrinsic phase, where the accuracy of model predictions is paramount for good performance. This is reflected in the downstream task performance evaluations in Fig. 8, where RaIR + CEE-US consistently outperforms both RaIR and CEE-US in the assembly tasks. Note that as regularity explicitly favours alignments such as stacks, unlike CEE-US, these dynamics are explored better, leading to higher success rates. This also showcases the importance of guiding free play towards regularity. In Fig. S5, the results for the pick & place, throwing and flipping tasks are shown. Due to the increased robustness of the model with the disagreement term, we indeed observe better performance for RaIR + CEE-US and CEE-US for the Pick & Place task compared to pure RaIR. This is also true for the throwing task. However, another contributing factor here is that models with disagreement favor more chaotic behaviors and perform more 'throwing'-like behaviors during free play. As CEE-US has no bias towards regularity, it performs best, whereas pure RaIR performs worse than RaIR + CEE-US. Interestingly, for the flipping 4 objects task we found performance for RaIR and RaIR + CEE-US to be comparable despite the significantly reduced amount of time spent flipping objects in the case of pure RaIR, as can be seen in Fig. 5. Upon inspecting the data generated during free-play, we hypothesize this is because unlike RaIR + CEE-US, which flips and rolls objects together in a chaotic way, we found RaIR to produce more isolated flipping of individual objects.

D Experiment Results for Re-creating Existing Patterns

As presented in Sec. 3.4, we test whether we can re-create existing regularities in the environment by simply optimizing for RaIR with iCEM, using ground truth models. As for the pyramid with

RaIR

RaIR+CEE-US

CEE-US

RND

Dis

success rate

0

100

200

300

Free Play Iteration

0%

50%

100%

(a) Pick & Place 6

0

100

200

300

Free Play Iteration

0%

50%

100%

(b) Throw 4

0

100

200

300

Free Play Iteration

0%

50%

100%

3 objects in Sec. 3.4, we initialize different regular structures outside of the robot's manipulability range and test whether these regular patterns can be re-created, merely by maximizing for RaIR. We test for the re-creation of a singletower with 3 and 4 objects, 2 towers with 2 objects each (referred to as multitower 2+2), as well as a spaced line and a rhombus with 4 objects. We test this with ground truth models for 15 independent rollouts for each structure and report the re-creation rates. Example rollouts are illustrated in Fig. S6. Note that due to the limited sample-budget with iCEM and the finite-horizon, we don't necessarily converge to the global minima, which corresponds to the full recreation of the structure. However, in all of the tested cases, the generated structures repeat at least one prominent sub-structure present in the underlying regular constellation by optimizing for RaIR.

For Singletower 3 , the entire stack of 3 gets recreated 73% of the time. A partial recreation with a stack of 2 blocks is observed in all but one of the remaining cases.

When a Singletower 4 is initialized outside of the robot's range, the full tower with 4 blocks gets recreated 40% of the time. In the remaining cases, either a tower of 3 (33%) or towers of 2 (27%) are built.

For the challenging Multitower 2+2 case, the two towers are built, with the same distance to each other as in the original pattern, 20% of the time. An example of this 'complete' recreation is illustrated in Fig. S6c. Otherwise, 53% of the time a stack of 2 is built (Fig. S6d) or a spaced line repeating the relative position of the two towers in the original pattern.

For the patterns on the ground, namely Spaced Line and Rhombus , the recreation rates are higher since the exploration problem is less prominent. At least 75% of the original pattern is re-created at each rollout, i.e. for the case of 4 objects, at least 3 objects follow the original pattern. The complete Spaced Line is recreated 80% and the entire rhombus 73% of the test rollouts.

In these experiments, we use RaIR with ϕ ( s i , s j ) = { ( |⌊ s i,x -s j,x ⌉| , |⌊ s i,y -s j,y ⌉| , |⌊ s i,z -s j,z ⌉| ) } . This is because in the case of existing structures in the scene, we don't need/want to inject any biases, e.g. towards vertical alignments, into the optimization. As the existing pattern is outside of the manipulability range of the robot, re-creating the pattern becomes a direct global optimum for RaIR, as all regularities reoccur. However, if we restrict ourselves to the x -y subspace, this is no longer the case: even for a rhombus, a stack built with the blocks in-reach becomes the global optimum. This is because all the blocks in the stack then have the same x -y relation to the blocks in the rhombus.

E Experiment Results for RaIR in Colored SHAPEGRIDWORLD with GT Models

In colored SHAPEGRIDWORLD, S obj includes not only the x -y positions, but also a binary variable encoding the different classes of colors. For c colors, we have an encoding of length ceil(log 2 c ) . So when we are computing relational RaIR, we also take into account the relations between colors. Take as example the case where we have two color classes blue (0) and orange (1) and use absolute relational ϕ : if all blue objects are in a horizontal spaced line, we have repetitions of | [ x i + δ x , 0 , 0] -[ x i , 0 , 0] | = [ δ x , 0 , 0] between neighboring objects with spacing δ x . If we have a line with an alternating pattern of colors in sequence, then we again have redundancies. For instance, for neighboring entities orange-blue we have: | [ x i + δ x , 0 , 1] -[ x i , 0 , 0] | = [ δ x , 0 , 1] . Since we have absolute relational ϕ , we get the same symbol for the blue-orange neighbors | [ x i + δ x , 0 , 0] -[ x i , 0 , 1] | = [ δ x , 0 , 1] .

Start

End

(a) Singletower 3

Start

End

(b) Singletower 4

Start

End

(c) Multitower 2+2

Start

End

(d) Multitower 2+2 (partial recreation)

Start

End

(e) Spaced line 4

Start

End

As a result, a blue-orange-blue-orange-. . . alternating line is also an optimum for RaIR. Additional patterns generated with RaIR in colored SHAPEGRIDWORLD with 2 and 3 colors using GT models are shown in Fig. S7.

F Experiment Results for RaIR in Custom CONSTRUCTION

We test regularity on a variant of CONSTRUCTION with diverse object shapes and masses (more details in Appendix. I.1).

F.1 RaIR with GT models

We apply RaIR to custom CONSTRUCTION with different shapes and masses, where we compute regularity between the x -y positions of the center of mass of objects. Here our aim is to show the generality of RaIR and that we are not constrained to identical entities. The generated constellations when optimizing for RaIR using GT models are showcased in Fig. S8.

F.2 RaIR in Free Play with Learned Models

We run free play in the custom CONSTRUCTION. During free play, we initialize one instance from each object type such that we have 1 cube, 1 ball, 1 column and 1 flat block. We then evaluate the zero-shot downstream task performance of RaIR + CEE-US and CEE-US for 3 different stacking tasks (Fig. S9). The results are shown in Table S2, where we obtain higher success rates with RaIR+CEE-US.

The free play parameters (model and controller settings) are identical to the standard CONSTRUCTION environment given in Table S6. The only difference is, we include a categorical variable to encode the different object types and append this to the observation vector for all objects. This is done in the same fashion as the static observations proposed in Sancaktar et al. [12]. For RaIR + CEE-US, we again use λ = 0 . 1 .

Stack Cube + Ball

Stack Cube + Flat

Block

Stack Cube +

Column + Ball

Table S2 : Zero-shot downstream task generalization performance of RaIR + CEE-US for custom CONSTRUCTION. Results are shown for five independent seeds, for models evaluated after 300 free play iterations (equivalent to 2K × 300 = 600 K transitions).

G Experiment Results for RaIR in the Quadruped and Walker Environments

We showcase that regularity also finds application in locomotion environments, where we compute RaIR on the world coordinates of the robot joints. Environment details for both Quadruped and Walker can be found in Appendix. I.1.

We run free play with RaIR + CEE-US ( λ = 0 . 02 ) and CEE-US in the Quadruped environment. With RaIR + CEE-US, we indeed observe more regular poses during free play, as reflected in the RaIR values shown in Fig. S10. We evaluate the zero-shot downstream performance of the learned models for the Roboyoga poses shown in Fig. S11. The success rates for the models at the end of free play (Table S3), as well as their temporal evolution (Fig. S12) are reported. We evaluate success for a pose as having a distance smaller than a threshold of 0.4 to the goal pose for more than 10 timesteps. We only see marginal improvements with RaIR + CEE-US compared to CEE-US. We hypothesize that this is because model-based approaches can generalize very well in locomotion environments and CEE-US also provides sufficient exploration. It

0

50

100

150

200

Free Play Iteration

-3.0

-2.6

-2.2

RaIR

only comes down to whether the pose can be held longer stably, where in the balancing poses RaIR+CEE-US has a slight edge over the course of free play (especially for Balance Front ). For any goal that is shown in Fig. S11 and not included in Table S3 (i.e. the poses with Quadruped lying on its back), we get perfect success rates for both RaIR + CEE-US and CEE-US. We didn't run free play in the Walker environment, but we expect model-based approaches to be very strong there altogether, as the environment dynamics are very simple.

Stand Leg Up

Attack

Balance Back

Balance Diagonal

Stand

Stand Rotated

Legs Up

Kneel

Side Angle

Stand Up

Lean Back

Boat

Bridge

One Foot Up

Arabesque

H CEE-US

In this section, we present the details of CEE-US [12], which we build upon in this work. CEE-US uses structured world models together with model-based planning during exploration, achieving increased sample-efficiency and superior downstream task performance compared to other intrinsicallymotivated RL baselines. The free-play pseudocode is presented in Alg. S1. This free-play structure is used for all methods presented in our paper by swapping out the intrinsic reward term (line 5)

RaIR+CEE-US

CEE-US

success rate

0

50

100

150

200

Free Play Iteration

0%

50%

100%

(a) Stand Leg Up

0

50

100

150

200

Free Play Iteration

0%

50%

100%

(b) Attack

0

50

100

150

200

Free Play Iteration

0%

50%

100%

(c) Balance Front

0

50

100

150

200

Free Play Iteration

0%

50%

100%

(d) Balance Back

success rate

0

50

100

150

200

Free Play Iteration

0%

50%

100%

(e) Balance Diagonal

0

50

100

150

200

Free Play Iteration

0%

50%

100%

0

50

100

150

200

Free Play Iteration

0%

50%

100%

(g) Lie Side Back

with only ensemble disagreement (CEE-US), combination of our regularity objective with ensemble disagreement (RaIR + CEE-US) or pure regularity (RaIR), as well as the other baselines RND and Dis (see Sec. J).

Algorithm S1 Free Play in Intrinsic Phase (taken from [12])

H.1 GNN Architectural Details

Message-passing Graph Neural Networks (GNN) are deployed as world models. The same GNN architecture is used as in CEE-US [12]. For these structured world models, we consider objectfactorized state spaces with S = ( S obj ) N ×S robot. Each node in the GNN corresponds to an object and the robot/actuated agent is differentiated from the object nodes as a global node. The concatenation of the robots's state s robot t and the action a t is represented as a global context c = [ s robot t , a t ] . We have a fully-connected GNN. The node update function g node and the edge update function g edge model the dynamics of the entities/objects, and their pairwise interactions respectively. These functions are both Multilayer Perceptrons (MLP). In the following, we denote the state of the i -th object s t, obj i at timestep t as s i t for simplicity. The object node attributes in the GNN are updated as:

̸

where [ · , . . . ] denotes concatenation, e ( i,j ) t is the edge attribute between two neighboring nodes ( i, j ) . For the permutation-invariant aggregation function given by aggr , we use the mean.

The robot state, which is treated as a global node, is computed using the global aggregation of all edges with a separate global node MLP g global :

Moreover, the GNN predicts the changes in the dynamics such that ˜ s t +1 = s t +GNN( s t , a t ) .

H.2 Planning Details

For planning, we use the improved Cross-Entropy Method (iCEM) [19]. The planner minimizes the cost, corresponding to negative reward c ( s t , a t , s t +1 ) = -r ( s t , a t , s t +1 ) , where r can be intrinsic rewards r intrinsic or extrinsic task rewards r task. The extrinsic task rewards are assumed to be given by the environment.

At each timestep t in the environment, the planner samples P action sequences, each with length H , i.e. the planning horizon. These actions are rolled out either in the ground truth model (perfect simulations) or in the imagination of a learned model (imperfect simulations), generating corresponding P state sequences with length H . In order to assign a cost to each of the P trajectories, we need to aggregate the cost over the horizon H . A typical choice here is sum , where the cost over the length of the trajectory is simply summed up: cost ( p ) = ∑ H -1 h =0 c ( s ( p ) t + h , a ( p ) t + h , s ( p ) t + h +1 ) .

However, this type of aggregation is not suitable for cases where a decrease in cost can in general be preceded by an initial increase. In these cases, using the mode best , that assigns the plan p the cost of the 'best' timestep over the planning horizon with cost ( p ) = min ( { c ( s ( p ) t + h , a ( p ) t + h , s ( p ) t + h +1 ) } H -1 h =0 ) is a better suited choice. We also empirically found this controller mode to be better at picking up sparse signals. What we mean here is that, in the example of stacking, it is hard to find a sampled trajectory that stacks the objects in a stable way with a limited sample-budget as this poses an exploration challenge. However, if we manage to find an action sequence that brings the cubes on top of each other, albeit in an unstable way, favoring this solution with best and keeping this solution in the elite set is beneficial. This can be explained as follows: In iCEM the K plans with the lowest assigned cost are chosen to be the elite set, which is then used to fit the sampling distribution of iCEM. As a fraction ξ of these elites is potentially shifted to the next internal iCEM iteration ( keep_elites ), and possibly to the next timestep ( shift_elites ), keeping these solutions that 'fail' and yet bring us closer to the actual solution provides a better strategy to solve tasks which pose an exploration challenge such as stacking. Here, we are also relying on the fact that we are re-running optimization every timestep t in the environment with online model predictive control, such that we have the opportunity to correct these initially 'wrong' solutions and find their 'stable' counterparts. Note that this mode of the controller is a more unstable mode compared to sum . Especially with imperfect world models, where the model can hallucinate as the model errors accumulate over the planning horizon, mode best can pick up these falsely imagined future states with low cost. It also doesn't account for the fact that the planned trajectory keeps the lowest cost over multiple timesteps, such that a trajectory where an object flies through the goal location for a single timestep has the same cost as a trajectory where the object lands in the goal position and stays there. To account for this, we use cost ( p ) = min ( { c ( s ( p ) t + h , a ( p ) t + h , s ( p ) t + h +1 ) } H -1 h =1 ) , where we don't take into account the first timestep of the plan with h = 0 . Although this is not a robust solution, we found it to empirically work well. Quantitatively, stacking 3 objects when planning with ground truth models yields 99% success rate for mode best , whereas only 47 % success rate for sum , using the same reward function in both cases. Even in the case of perfect dynamics predictions with GT models, this showcases the importance of the controller mode to be able solve tasks with sparse reward signals. This is further showcased in Table S1, when comparing the results for CEE-US with default best and sum .

I Experiment Details

In this section, we provide experimental details and hyperparameter settings.

I.1 Environment Details

SHAPEGRIDWORLD This is a discrete 2D grid, where each entity/agent is controlled individually in the x -y directions. This means entities are controlled one at a time and actuation keeps iterating over the entities, where we use an object persistency of 10 timesteps. The action a t is 2 dimensional, controlling the agent in x -y directions separately and is applied on the current actuated entity i in the grid. As we are operating in a discrete grid, the actions are actually discrete such that the agent can move one grid cell to the left/right and up/right (if the target grid cell is not occupied) or stay at the current grid cell. In order to make this environment work with the default iCEM implementation with a Gaussian sampling distribution, we perform a discretization step before inputting the sampled actions to the environment. For the experiment results with GT models presented in Fig. 1, we use a grid size of 25 × 25 with 16 and 32 objects. In colored SHAPEGRIDWORLD with c colors, each object is assigned one of the c available color options, and the observation for each entity also includes binary encoding for color with length ceil(log 2 c ) .

CONSTRUCTION This is a multi-object manipulation environment as an extension of the Fetch Pick & Place environment proposed in [24]. We also applied the two modifications from Sancaktar et al. [12]. 1) The table in front of the robot is replaced with a large plane such that objects cannot fall off during free play, but can still be thrown/pushed outside of the robot's reach. 2) In Li et al. [24], the object state also contained the object's position relative to the gripper which was removed, as it already introduces a relational bias in the raw state representation. Details on the dimensionalities of the object and robot state spaces can be found in Table S5.

Custom CONSTRUCTION As an extension of the standard CONSTRUCTION, there are 4 different object types with different masses: cube (same as in original CONSTRUCTION) with size 5 cm and mass 2 (mass in the default unit of Mujoco), flat block with size 16 × 7 × 3 cm and mass 1.5, column (upright block) with size 3 × 3 × 10 cm and mass 1, and a ball with diameter 8cm and mass 1. We also include categorical variables for each object type as static observations. During free play, we initialize one of each object type, so we have 4 objects. Otherwise the specifications are the same as in Table S5.

Quadruped & Walker These environments are taken from the DeepMind Control Suite [45]. To be able to perform RaIR computations, we modify the observation space of each environment and append the world coordinates of the robot's joints to the observations vector. In Quadruped, we take the toe and knee positions. We used x -y -z coordinates with the GT models (results shown in Fig. 10), and x -y for free play in Quadruped with learned models. In principle, we can also include the ankle and hip joints, but we observed similar performance in this case. For Walker, we append the feet, leg and torso x -y positions.

I.2 Parameters for Ground Truth Model Experiments with RaIR

The controller parameters used when optimizing RaIR with ground truth (GT) models are given in Table S4. To compute RaIR, we perform a discretization step in the CONSTRUCTION environment as it is continuous. For GT models, that produce perfect mental simulations, we can choose a small bin size of 1cm. In comparison, the size of one block in the environment is 5 cm. The bin size also gives us the upper bound of the regularity precision that can be achieved during optimization, e.g. a perfectly aligned stack vs. a zigzagged stack. Note however that the higher the precision is, the harder it typically gets for the controller to converge to global optima with a horizon of 30 timesteps and a limited sample-budget.

As discussed in Sec. A, this is also a constraint when we are computing RaIR in the x -y -z subspace. Due to this, for the re-creation of existing patterns experiments presented in Sec. 3.4 and Sec. D, we compute RaIR with a bin size of 2.5cm and increase the number of sampled trajectories P to 512. Although we could further increase the bin size, we choose this value to not negatively impact the precision of the re-created structures. In custom CONSTRUCTION, we also use a bin size of 2 . 5 cm, when optimizing RaIR with GT models.

For Quadruped, we also use a bin size of 5cm. For the GT model experiments, we compute RaIR over the x -y (-z ) positions of the Quadruped's toe and knee joints. We didn't see significant qualitative differences when including the z -positions in the regularity computation. For Walker, we use a bin size of 1cm and compute RaIR over the x -y positions of the Walker's feet, legs and torso.

I.3 Free Play with Learned Models

The environment properties with the episode lengths and model training frequencies are given in Table S5. Six objects are present in CONSTRUCTION during free play. The parameters for the GNN model architecture as well as the training parameters for model learning are listed in Table S6. For the RaIR computations in free play, we use a bin size of 5cm, which is equivalent to the size of a block.

For custom CONSTRUCTION, 4 objects (one of each object type) are present during free play. The hyperparameters are the same as for standard CONSTRUCTION, and a bin size of 5cm is used.

The set of the hyperparameters for the iCEM controller used in the intrinsic phase of RaIR + CEE-US, RaIR and CEE-US are the same as presented in Table S4. The only difference to the GT model case

(a) General settings.

(b) Environment-specific settings.

is, we use a planning horizon of 20 timesteps for free play. We also use the exact same settings for the RND baseline. For Disagreement, we have a planning horizon of 1 during free play.

For the Quadruped environment, we use an MLP ensemble instead of a GNN ensemble. The hyperparameters are reported in Table S7. The iCEM controller parameters used in free play for Quadruped are the same as the ones for GT models in Table S4, except for the noise strength σ init = 0 . 5 .

I.4 Extrinsic Phase: Zero-shot Downstream Task Generalization

In this section, we provide details on the extrinsic phase following free play, where the learned GNN ensemble is used to solve downstream tasks zero-shot via model-based planning.

I.4.1 Details on Downstream Tasks and Reward Functions

The reward functions for all the CONSTRUCTION downstream tasks are computed as specified in Sancaktar et al. [12], where for all the assembly tasks we use the same structure as in the stacking reward. However, we do one modification to the original reward computation in the assembly tasks. The assembly task reward is sparse incremental with reward shaping, where the reward also contains the distance between the gripper and the position of the next block to be stacked. We modify how the next block ID is computed in the original implementation from Sancaktar et al. [12]. Instead of naively checking the number of unsolved objects to obtain the next block ID irrespective of order, we determine the next block to be the next unsolved block in the order. We found this modification to be important especially for the Pyramid tasks, where the sub-optimal next block computation might lead to the agent receiving a reward to be close to the wrong block, in the case the robot places blocks with i > next_block_id to their goal locations with just the sparse reward.

For the RoboYoga experiments, the reward function is the same as proposed in [25] (taken from their code in https://github.com/orybkin/lexa-benchmark ). It is a dense reward as the shortest angle distance to the goal pose.

I.4.2 Planning Details for Downstream Tasks

The controller settings for the different downstream tasks are shown in Table S8, which are the same settings used in [12].

J Details on Baselines

J.1 Random Network Distillation

Random network distillation (RND) [21] has two networks: a predictor network and a randomlyinitialized and fixed target network. The predictor network f θ : S - → R k is trained on the data collected by the agent to match the outputs of the target network ˆ f : S - → R k , where k is the embedding dimension. The error between the predictor and target networks for the state s t is used as the intrinsic reward signal:

where the idea is that this error will be higher for novel states. RND can be seen as an extension of count-based methods to continuous domains. The parameters used for the RND module is given in Table S9.

Unlike prior work [21], we don't use RND with an exploration policy. Instead, we have a custom implementation using the same model-based planning backbone for RND as CEE-US, RaIR + CEE-US and RaIR. By using RND with iCEM instead of an exploration policy, we gain sample-efficiency, as also showcased in [12]. Note that although RND intrinsic reward itself is detached from the model and does not require an ensemble, we still deploy an ensemble of GNNs in this case. In previous work [12], a single GNN model was used. In order to compute the RND intrinsic reward, we take the mean predictions of the ensembles. We do this because ensembles have been shown to be more robust models [18], such that there is no inherent disadvantage to the RND baseline.

J.2 Disagreement

The Disagreement baseline [10] uses the same intrinsic reward metric as CEE-US, given in equation (Eq. 5). We implement it with the same model-based planning backbone as in CEE-US, replacing the exploration policy in [10]. The only difference to CEE-US is that the planning horizon during free play is only 1-step. The model and iCEM controller settings are otherwise the same as for CEE-US, RaIR + CEE-US and RaIR. Essentially, the comparison between this baseline and CEE-US showcases the importance of multi-step planning for future novelty. Note that the planning horizon

is only 1 for the free play phase. For fair comparison, we use the same horizon when solving the downstream tasks as for all other methods.

K Discussion on Goal-conditioned RL for Unsupervised Exploration

In our setup with model-based planning for unsupervised RL, the focus is on distilling knowledge into a world model that is robust enough such that we can plan for diverse downstream tasks afterwards. In the case of the goal-conditioned RL literature, the focus is on distilling this knowledge into a goal-conditioned policy, as in [25, 37, 38].

Although the exploration itself is unsupervised and guided by intrinsic rewards, the moment we solve a downstream task we need a goal specification from the environment. Whether you use a goal-conditioned policy to solve it or as in our case, use this goal to compute a reward function, is a separate distinction.

In goal-conditioned RL, you also need access to rewards at training time: the goals are used to compute this reward signal that is used to train the policy and value function. In a way, GCRL is supervised in terms of how you specify your goals (which goal space you use) and the rewards with respect to these goals. We consider the reward function in GCRL to be actually known to the agent. The nice part in GCRL is, you can put in very general supervision with minimal assumptions, that don't require much knowledge of environment dynamics. However, doing so, GCRL is restricting the class of tasks you can solve, namely to tasks that can be described as a goal. But in a wide range of environments, it is hard to specify tasks as a goal: for example running as fast as possible in Halfcheetah.

Our goal is to learn a robust world model that we can use for planning in a later extrinsic phase. In our setup, during the intrinsic phase we make similar assumptions as Plan2Explore [11]. For the extrinsic phase, we need some specification of the task. For CEE-US, a reward function is provided at test time, whereas for GCRL it is a goal (which in fact also implies a reward function).

In model-based planning, since we have complete freedom of choice for the reward function, we can specify any task that can be solved in an MDP setting. Note that we can use our setup to also optimize a goal-conditioned sparse reward. The challenge for us is the limitedness of the zero-order trajectory optimizer: solving long horizon tasks with just sparse rewards becomes challenging, especially with imperfect world models. That's why in practice, we fall back to dense rewards for planning.

L Connections between Compression and RaIR

Our regularity objective, that seeks out low-entropy states with high redundancy, shares close ties with compression, and specifically with lossless compression using entropy coding.

We implemented a version of our regularity idea using the lossless compression algorithm bzip2 corresponding to the direct version of RaIR with order k = 1 (Sec. 2.2). In this case, we describe the state s directly by the properties of each of the entities and the function ϕ : S obj →{X} + , that maps each entity to a set of symbols and obtain Φ( s ) = ⋓ N i =1 ϕ ( s obj ,i ) as a union of all symbols for N objects. Instead of computing the entropy for the frequencies of occurrence in the resulting multiset of symbols like in RaIR, we instead transform these symbols into bytes and compress them with bzip2 . We then define the intrinsic reward for compression as the negative length of the compressed ByteString such that:

We also managed to create regular shapes and patterns when optimizing for r compression via planning with ground truth models and also for free play with learned models. The reason we chose not to pursue this direction was because 1) lossless compression algorithms like bzip2 don't perform as well on short ByteStrings, which is the case for us, as e.g. in CONSTRUCTION, we compress only 6 objects with their corresponding x -y positions, 2) artifacts are introducted to the regularity/compression metric by the transformation into bytes, where certain symbols become more compressible than others in this representation without any added regularity. As a result, we preferred our formulation with RaIR as it provides better control over the generated patterns and structures.

M Discussion on Object-centric Representations

Wewant to further discuss our reliance on object-centric or more broadly entity-centric representations of the world for our regularity computations. We view the world as a collection of entities and their interactions. This is a very general assumption, and one that is commonly used for its potential to improve sample efficiency and generalization of learning algorithms in a wide range of domains [46-48]. The principles underlying entity segmentation and their importance for our perception are also studied in cognitive science [49-51].

Entity itself is a very abstract and fluid concept. For us as humans, it dynamically changes depending on the granularity of the control problem we want to solve. For example, if we want to move a plate of cookies from the counter to the table, each cookie on the plate is not treated as a separate entity, instead we view the whole plate as one. If we later want to eat the cookies, then the entityview changes: now each cookie is its own entity. (Example adapted from a private discussion with Maximilian Seitzer.)

Deciding on which level of abstraction to use in which control setting is non-trivial and a research direction which we find very exciting. In object-manipulation environments that are currently encountered in RL (with hard rigid bodies and non-composite structures), we don't find much ambiguity on entity identifications, hence our statement in Sec. 6 that observations per object in these cases are naturally disentangled. However, in open-ended real world scenarios it is indeed needed to dynamically choose the right level of abstraction in the perception of a scene and identify the necessary set of entities for the control problem at hand. Using RaIR for a real-world e.g. cluttered scene would be very challenging without closing this action-perception loop. With orthogonal research in visual perception, we are hopeful that this will be possible and we see these synergies as exciting future work.

N Code and Compute

Code is available on the project webpage https://sites.google.com/view/rair-project . We run the ground truth model experiments on CPUs. As we are using the true environment simulator as a model, each imagination step in the planning horizon takes as long as an environment step. We parallelize the ground truth models on 16 virtual cores The controller frequency in this case is ca. 0 . 25 Hz , for the settings given in Table S4. For the free-play phase, we have a fixed number of transitions collected at each training iteration, which get added to the replay buffer. After the data collection, the model is trained on the whole replay buffer for 25 epochs. Since the buffer size increases at each training iteration with newly collected data, for this fixed number of epochs, the corresponding number of model training updates and thus also the runtime of the iteration, increase throughout free play. For RaIR + CEE-US, the full free-play (300 training iterations) in CONSTRUCTION with 6 objects, where overall 600K data points are collected, takes roughly 87 hours using a single GPU (NVIDIA GeForce RTX 3060) and 6 cores on an AMD Ryzen 9 5900X Processor. The majority of this time is spent on the model training after data collection. The controller frequency for the collected rollouts with RaIR and the epistemic uncertainty calculations using a GNN ensemble is ca. 5 Hz .

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis work proposes an exploration approach that prioritizes \"regularity\" in the exploration behavior, in contrast to typical novelty-seeking exploration. The objective is defined as the minimization of entropy over user-defined, object-centric features. The method is primarily evaluated on block rearrangement in the robotic manipulation setting.\n\n---\n**8/22/23 Update after author response**\n\nMy major concern with this paper was the generality of RaIR. In my opinion, there are two ways to create a good exploration method. The first is to come up with a general method with little-to-no assumptions that works across a variety of tasks. The second is to create a specialized method, be upfront about the assumptions made, and get very good performance by leveraging those assumptions. \n\nBefore the rebuttal, RaIR was introduced as a general exploration approach, yet its experiments were limited to stacking objects of identical shapes and sizes. While stacking is a very hard exploration challenge, I expected to see environments beyond object manipulation like locomotion, video games, etc. Furthermore, RaIR does depend on the key assumption of object-centric state representation, which it uses in both its objective and dynamics model. \n\nThe authors provided some followup experiments in the rebuttal. They evaluated their method on a more generalized version of their construction environment, as well as on a Quadruped locomotion environment. They also provided some additional new environments (RoboDesk, Walker),  but did not run exploration experiments on them, only sanity checks. \n\nSo to conclude:\n\n**My concern on generality of the method:** With the new exploration experiments on generalized construction and quadruped, the authors have shown that RaIR can handle diverse objects, and can work in a single non-manipulation task. Therefore my concern is somewhat alleviated. I would recommend the authors to show more diverse environments in general though, to remove any remaining doubt from future readers that RaIR is not engineered towards a particular type of task.  \n\n**My other concerns about object-centric assumption, baselines, related work, novelty, etc:** The authors have done a good job on addressing my concerns here. I hope these discussions will go into the next version of the paper, and influence how RaIR is presented.\n\nAs of such, I will raise my score to 5. The authors have passed the threshold for addressing my main concern on generality, given the limited rebuttal time.\n\nFor future improvements, I would recommend the following:\n1) Continue adding more diverse environments to show RaIR is a general exploration method.\n2) Study the assumption of object-centric representation a bit more - what happens if we use RaIR with an unstructured dynamics model? Will the regularity objective still work?\n\nSTRENGTHS:\n- The high level idea of using regularity to focus exploration is well motivated, and addresses the weaknesses of the more common novelty-seeking exploration methods. \n- This method explores well in block stacking / rearrangement, a hard exploration problem. The experimental section analyzes the block stacking task extensively.\n- The manuscript is generally well written.\n\nWEAKNESSES:\nMy main concern is over the scope and generality of this method. While the introduction motivates regularity as a general-purpose exploration objective, the actual implementation and assumptions made to implement the objective restricts it to a very particular set of tasks - manipulation of identical and well characterized objects.  I would have expected, from the introduction, a method that could apply to all sorts of domains (e.g. robotic locomotion, video games) instead of only object manipulation tasks.\n\nIn RaIR's environments, all objects are the same shape and have the same physics. For the construction environment, the authors only apply RaIR to the x-y state space to bias towards vertical alignments. However, this x-y assumption would break if the objects were not identical, since the vertical alignment would now be object dependent. \n\nAnother major limitation (which I appreciate the authors mentioning in Sec. 5)  seems to be the assumption of a known state space and dependence on user features. RaIR seems to hinge on having a state space that is interpretable, so that the user can provide user-defined features to bias the exploration.  I am concerned about the scalability of providing user defined features, particularly if the state space becomes more complex, or uninterpretable (e.g. pretrained ImageNet representations). However, if RaIR claims to be a general exploration method, it should address such problems. \n\nFinally, the experiment section is lacking baselines. There is only 1 novelty-seeking baseline that the authors compare against, which is by design going to be suboptimal in assembly tasks. I would like to see how RaIR compares to an exploration baseline that actually seeks regularity in its objective as well. I believe SMiRL (1) to be a viable baseline, as it also seeks to minimize surprise. I would also like to see some baselines from the unsupervised goal-conditioned RL literature, as those methods also do stacking. \n\n## Moderate concerns:\n- Claims of Novelty - The authors claim in the first line of the abstract that they propose regularity as a novel reward signal for intrinsic motivation. However, the high level idea of regularity / stability / niche-seeking has been explored as an intrinsic motivation signal for RL(1), and more generally in the active inference literature (2). I would like to authors to revise their claims of novelty, and include the relevant literature.\n- Related work is rather short. Only a small paragraph is dedicated to intrinsic motivation. A rather long paragraph is dedicated to Compression, which is not necessary.  I would recommend the following changes.\n\t- Expand on intrinsic motivation in RL, particularly methods that also do manipulation tasks and how RaIR compares to them. Mention (1,2) as well, as they also explicitly mention concepts similar to regularity (stability, niche-seeking). Another omission is the unsupervised skill discovery literature, where some works (3) also show manipulation tasks.  \n\t- Paragraph on compression is unnecessarily long and detailed - it is sufficient to briefly mention compression progress and move the rest to the appendix if necessary. That space should be used mainly on Intrinsic Motivation.\n\n## Minor concerns:\n- Line 84: what is $n_a, n_s$?\n- Line 124: I found the \"Relational RaIR of order K\" to be a bit abstract, and I needed to read it a few times over to fully understand it. Having a figure like Figure 3 for this section would have been very helpful.\n- Line 220: Please point to a figure.\n- Recreation experiment - while the recreation behavior is interesting, the authors do not include any practical applications or downstream implications for this phenomenon. I would prefer this section be moved to the appendix and the space be used for other things (e.g. experiment in new domain) or extended related work.\n\t\n1. Berseth et al, SMiRL: Surprise Minimizing Reinforcement Learning in Unstable Environments, ICLR 2021\n2. Friston et al, Active inference and learning. Neuroscience & Biobehavioral Reviews 2016\n3. Zhao et al,  Mutual Information State Intrinsic Control, ICLR 2021\n\nQUESTIONS:\nWhat is the scope of RaIR? I can see two ways to present RaIR. \n\nOne way, and the way the paper is written, is to promise a broad scope. However, can RaIR be applied to more complex and realistic settings both in the manipulation domain and beyond it?  I would suggest for experiments that show RaIR can be applied to other domains, and in more complicated scenarios.\n\nThe other way, is to narrow down the scope and claims of RaIR to an intrinsic reward for object manipulation, which itself is a very challenging and interesting domain. I would like to see experiments in more realistic manipulation setups (different shapes, sizes, masses, etc.), as mentioned above.\n\nLIMITATIONS:\nThe authors did address a main limitation, which is the assumption to an object centric, interpretable state space. \nAnother limitation not mentioned is the need to define task-specific features for RaIR.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 5, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-06-11"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper presents a new approach that utilizes a regularity as an intrinsic reward, which aims to regularize, or constrain, the search space of exploration towards state regions where the entities are more likely to have structured patterns. To this end, the paper first formulates how to define regularity as an intrinsic reward and investigates several candidates that can be used as a regularity. And this intrinsic reward is combined with the disagreement-based intrinsic rewards to guide the exploration. Experiments show that optimizing this intrinsic reward (by using the ground-truth models) indeed leads to behaviors that induces structured patterns in the environments, and further show that this can be useful in case everything is learned end-to-end.\n\nSTRENGTHS:\n- Interesting formulation based on a good intuition. I enjoyed reading the paper.\n- Clear writing helpful for understanding the method. This should be highlighted as understanding the concept would have been difficult without the clear writing.\n- Experiments are conducted to support the claims by using ground-truth models and learned models. It starts from showing the potential benefit of using the proposed reward by using the ground-truth model, then shows that it can be useful with learned models, and finally shows that it can be useful for improving the performance on downstream tasks.\n\nWEAKNESSES:\n- In contrast to clear writing until the Method section, Experiments section is a bit difficult to parse, especially it is difficult to understand the metrics used in Section 3.2. Even though they are explained in the text, further emphasis and clear description on the meaning of metrics (e.g., what relative time means, how to interpret objects in the air, flipped) could be helpful.\n- As the authors already highlighted in the Limitation section, this method severely depends on the accessiblity of ground-truth internal states of all the entities in the environment, which is not feasible in practice. But this is understandable as learning such information is not the main focus of the paper.\n- Another limitation of this paper is the difficulty of balancing (i) the chaotic exploration from novelty-seeking intrinsic rewards. (in this paper, disagreement-based reward) and (ii) regularity-seeking intrinsic rewards. Of course there could be cases when injecting our prior of prefereing the structuredness can be helpful, but sometimes it's not, as can be seen in Throw and Flip experiments in Table 2. Would there be a way to automatically balance these two rewards, or in general, how can we control the magnitude of injecting our prior to the exploration process?\n- Experiments are a bit limited as only specific type of structuredness is investigated (for simplicity).\n\nQUESTIONS:\nI don't have further question other than the ones described in the Weaknesses.\n\nLIMITATIONS:\nYes.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-05"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper draws inspiration from a child’s development cycle and proposes the use of regularity as an intrinsic reward to help guide exploration in RL. It shows that existing metrics of state novelty can be combined with the proposed regularity metric to enable construction of regular structure during free play in a sample efficient manner. The proposed method is validated on experiments from two different tasks.\n\nSTRENGTHS:\n- The paper proposes an interesting and novel formulation of using regularity, in the form of symmetry, as an intrinsic reward to guide exploration during RL.\n- The paper builds on an existing novelty-based approach that uses the epistemic uncertainty obtained from an ensemble of models to guide exploration. The authors show that RaIR can be combined with such novelty-based rewards to further improve exploration in tasks that demand regularity. However, the authors acknowledge that for tasks that are chaotic like throwing and flipping, imposing such regularity hurts the performance (Table 2).\n- The paper shows impressive results for stacking towers in various configurations solely from intrinsic rewards while prior work achieves this using carefully constructed reward functions.\n- The paper demonstrates the effectiveness of pretraining with the regularity-based intrinsic reward in performing zero-shot generalization to downstream tasks (Sec. 3.3).\n- The authors also show a way of prompting (Sec. 3.4)  the agent to build certain structures by incentivizing it to raise regularity by building the same structure as the prompt.\n\nWEAKNESSES:\n- The authors organize the state in the form of graph nodes with each node representing a specific feature of an object or the agent. I had a few questions about this - \n  - How did the authors decide on this specific state representation? Does it have any advantages that a particularly attractive?\n  - It would be great if the authors could clarify how the graph is built (is it a fully-connected graph with all nodes)?\n   - The current form of state representation results in the authors using GNNs. However, the same representation can be dealt with using transformers. It would be interesting to add a transformer variant of RaIR.\n  - This state representation might make it difficult to apply the method to the real world where privileged state information is not available. I would be curious to hear the author’s thoughts on this. The limitations section mentions that the same operation could be done on a latent space but it's unclear if the RaIR would be applicable to such a latent space (which might not be as good as the privileged information) based state representation.\n- Since the method uses iCEM for policy learning, it seems like such a method of optimizing the policy for each new configuration won’t scale with reinforcement learning. It would be great if the authors could comment on this.\n- How many training steps does it take to train the transition model? This would throw some light on the sample-efficiency of the method.\nWhat is “relative time” in Figure 5?\n- In Sec. 3.4, why does the agent not arrange everything in a single line to maximize regularity? Is it because of the finite horizon optimization (meaning that tasks that need more time steps would be preferred less as compared to tasks giving higher immediate rewards)?\n\nQUESTIONS:\nIt would be great if the authors could address the comments mentioned in the “Weaknesses” section.\n\nLIMITATIONS:\nThe paper does have a limitations section. The authors must also acknowledge the following from the “Weaknesses” section - \n\n```This state representation might make it difficult to apply the method to the real world where privileged state information is not available. I would be curious to hear the author’s thoughts on this. The limitations section mentions that the same operation could be done on a latent space but its unclear if the RaIR would be applicable to such a latent space (which might not be as good as the privileged information) based state representation.```", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-05"}

comments: SUMMARY:
Inspired by the fact that humans prefer symmetric patters, the author proposed to use regularity as intrinsic reward to encourage an agent to discover interesting environmental patterns. One example of regularity is to put boxes in a symmetric pattern. The results showed that this kind of intrinsic reward allows an RL agent to generate more interesting patterns in the absence of task rewards and also allows for better world model learning for zero-shot generalization.

STRENGTHS:
- The proposed method is novel and intuitive.
- The experiments are comprehensive, including both analysis of the emergent behaviors and the performance of applying to downstream tasks.

WEAKNESSES:
- Lack of baselines: For section 3.3, I believe RND and ICM can be a reasonable baselines to strengthen the significance of the proposed method.  It seems to me that the goal of section 3.3 is to showcase the practical benefit of RaIR in solving tasks (maximizing extrinsic rewards). To show practical benefit, it would be necessary to answer why the proposed method is a better choice than prior works. Comparison with RND and ICM would answer this question.
- Regularity considered in this paper are not well motivated. I agree the motivation of using regularity as intrinsic rewards is clear, but the specific design of each type of regularity in Table 1 are unclear to me. For example, why we consider these symmetry operations and why design distance in these ways.

QUESTIONS:
(See weakness)

LIMITATIONS:
Yes, the limitation is mentioned.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 2
DATE: 2023-07-08