COBRA: Data-Efficient Model-Based RL through Unsupervised Object Discovery and Curiosity-Driven Exploration

1. Introduction

• Data efficiency and robustness to new contexts remain persistent challenges for deep RL, particularly when agents must learn practical tasks with limited supervision.Source: §1 · p. 1
• COBRA combines object-based representations, curiosity-driven exploration, and model-based RL to pursue data efficiency and policy robustness.Source: §1 · p. 1
• Its technical contributions include unsupervised action-conditioned dynamics over slot-structured representations learned from raw pixels, a sampler for continuous actions, and an integrated agent architecture.Source: §1 · p. 1
• COBRA first learns to see, predict, and act during reward-free exploration, then uses those capacities to learn tasks through model-based RL.Source: §1 · p. 1

Curiosity-driven exploration

• Prior curiosity-driven exploration work introduced world models for planning or exploration, while later studies examined task-free exploration and intrinsic motivation; connections to object representations and policy robustness remained unexplored.Source: §2 · p. 2

Object-Oriented RL

• Object-oriented RL promises data efficiency and generalization through object and relation representations, but most existing approaches require hand-crafted representations or relations [Garnelo et al., 2016; Lázaro-Gredilla et al., 2019].Source: §2 · p. 2
• COBRA differs from prior object-oriented approaches by automatically discovering object representations from data without hand-crafting or supervision.Source: §2 · p. 2

Model-based RL

• Learning environmental models has advanced, but model-free alternatives remain difficult to beat asymptotically; concurrent work combined intrinsic motivation and model-based RL on continuous control tasks (Hafner et al., 2018).Source: §2 · p. 2
• COBRA learns a more structured latent representation by decomposing scenes into objects, which the authors hypothesize may help scale to more complex situations; they note their current policy-learning approach is more limited and comparison with Hafner et al. (2018) remains future work.Source: §2 · p. 2

3. The COBRA model

• The model has a vision module, transition model, exploration policy, and reward predictor, with the first three trained without reward and the reward predictor trained during the task phase.Source: §3 · p. 2, Figure 1 · p. 3

3.1. Unsupervised exploration phase

• During reward-free exploration, COBRA learns representations of objects, dynamics, and its action space from episodes initialized with one to seven randomly configured objects.Source: §3.1 · p. 3
• MONet decomposes scenes into object slots with shared latent semantics that capture properties such as position, angle, color, and shape, with unused slots able to represent no object.The paper reports that MONet accurately represents Spriteworld scenes in terms of objects.Source: §3.1 · p. 3
• The transition model applies a shared action-conditioned MLP slot-wise and trains predictions through decoded images and pixel loss, avoiding direct matching of potentially reordered slots.An auxiliary network predicts transition error and supplies the curiosity measure used to train exploration.Source: §3.1 · p. 4
• Because random actions rarely click on objects occupying about 1.7% of the image each, COBRA adversarially trains its exploration policy to seek actions with high predicted transition error.The authors describe this joint learning as a curiosity-driven cycle in which hard-to-predict actions tend to move objects.Source: §3.1 · pp. 4–5

3.2. Task phase

• The task phase freezes the vision module, transition model, and exploration policy, learns a separate reward predictor for each task, and selects actions through one-step model-based search.The agent samples candidate actions, predicts next states and rewards, then chooses the best action with epsilon-greedy selection.Source: §3.2 · p. 6

Environment

• Spriteworld is a 2D arena of colored sprites with occlusion but no collisions, using continuous four-dimensional click-and-push actions that move an object only when the position click lands inside it.Source: §4 · p. 6
• The environment supports compositionality, task-irrelevant perturbations, and combinatorial generalization without providing privileged object representations or object-specific actions.Source: §4 · p. 6

Tasks

• The six tasks span goal-finding, sorting, and clustering, with held-out tests varying target or distractor counts, task-irrelevant features, object combinations, or colors.Source: §4 · p. 7

Evaluation

• The data-efficiency measure is the fewest on-task environment steps needed to sustain 90% average performance for 30 consecutive episodes; gray bars mark agents that did not reach that threshold.Source: Figure 4 · p. 8

5. Results

• The baselines are MPO raw, which is model-free, and MPO handheld, which reuses COBRA’s vision module and exploration policy.MPO raw must be retrained for each task, whereas MPO handheld reuses the same pretraining as COBRA.Source: §5 · p. 8
• COBRA solves most tasks with 100–1,000 rewarded environment steps, compared with 10^6–10^8 for all MPO baselines; MPO raw must be entirely retrained across tasks.The unsupervised exploration phase is trained once and reused for any task. MPO handheld reuses the same amount of pre-training, making that data-efficiency comparison fairer.Source: §5 · p. 8
• COBRA far exceeds baselines on many robustness tests; the authors hypothesize that learning a reward function rather than a policy network reduces overfitting.They conjecture that learning goals and planning with a model can yield more robustness than model-free RL.Source: §5 · pp. 8–9
• In Figure 5, COBRA handles an added distractor, a held-out blue/red sorting pair, and clustering red/yellow objects after training on green/blue.The displayed target locations are for visualization only and are not provided to the agent.Source: Figure 5 · p. 9

6. Discussion

• COBRA combines unsupervised object-centric representations, curiosity-driven exploration, and model-based RL, with high data efficiency and robustness to task-irrelevant perturbations demonstrated in experiments.The authors describe this as, to their knowledge, the first architecture to bring these components together.Source: §6 · p. 9
• Spriteworld uses continuous actions without hand-crafted object definitions; discretizing actions proved difficult to train, while object-indexed actions would bypass object discovery.The authors also favor the continuous action space for its more realistic geometric relationship to the environment.Source: §6 · pp. 9–10
• The transition model omits object interactions and memory, so it cannot retain out-of-view objects in settings that require memory.The authors consider the lack of memory acceptable in fully observable Spriteworld and suggest GraphNet as a possible way to address interactions.Source: §6 · p. 10
• The deformation-based exploration policy learns basins around high-curiosity actions without equal-volume pressure, favoring objects in open space over those near boundaries or other objects.This did not affect performance in Spriteworld but might matter in more cluttered environments.Source: §6 · p. 10
• Transition-model error is unsuitable as a curiosity measure in unpredictable-dynamics environments; the authors suggest the Intrinsic Curiosity Module as a complementary alternative.They also note that goal parameterization is challenging when the number of objects varies.Source: §6 · p. 10
• One-step search is sufficient for the tested tasks, but complicated tasks may need multi-step rollouts; stateless planning cannot handle tasks requiring memory.A task where a goal location briefly flashes is given as an example of the memory boundary.Source: §6 · p. 10
• Future directions include richer environments and avatars, learning to set goals, multi-step exploration planning, and learning a task-time policy rather than using simple search.Source: §6 · p. 10

Supplementary material

• The supplementary material provides links to training and robustness-test videos for all tasks, plus a README summarizing them.Source: A. Agent Videos · p. 14

B.1. Vision Module

• MONet provides the unsupervised vision module; the implementation uses eight slots and a three-layer Spatial Broadcast Decoder, with color-channel rescaling to aid training.Eight slots exceed the environment’s maximum of six objects because MONet uses a background slot; the decoder and color settings were reported to improve decomposition slightly.Source: §B.1 · p. 14

B.2. Transition Module

• The transition model predicts the next scene and its own error from slot-wise inputs, using pixel reconstruction, error-prediction, and regularization losses summed without reweighting.The predicted error serves as a more efficient and stable proxy for transition error when training the exploration policy adversarially.Source: §B.2 · p. 14

B.3. Exploration Policy

• The exploration policy samples deformations from a four-dimensional Gaussian and is trained to maximize predicted transition error while regularizing deformation size and small scales.The policy adds its deformation to a uniformly sampled action; using a distribution rather than deterministic deformations improved training efficiency and stability.Source: §B.3 · p. 15

B.4. Reward Predictor

• The reward predictor uses a Relation Network with slot-pair and global MLPs to map the slot-structured state to a scalar reward prediction.The slot-pair MLP has layer sizes (128, 128), and the global MLP has sizes (128, 1).Source: §B.4 · p. 15

B.5. Baselines

• MPO baselines use either slot-structured or image-based actor–critic networks, and the reported layer-size sweeps found no overall performance gain from doubling MLP widths or changing layer counts.MPO training used trajectories of length 2; the hyperparameters are listed in Table S1.Source: §B.5 · p. 15, Table S1 · p. 16

C. Training details

• The vision module was pretrained on two million random static frames, and the exploration phase used at most 10^7 environment steps and could be reused across tasks.Task-phase COBRA training used 1,000 episodes, corresponding to 5,000–30,000 environment steps; MPO baselines also required the exploration phase for vision and exploration-policy training.Source: C. Training details · p. 16
• Actor-count effects complicate data-efficiency comparisons: fewer actors significantly degraded MPO learning efficiency, while the broader performance–data-efficiency trade-off was outside the study’s scope.The reported data-efficiency numbers should be treated cautiously, although the authors characterize the difference from baselines as several orders of magnitude.Source: C. Training details · pp. 16–17

D. Model Variations Considered

• The appendix presents model variations explored before settling on the transition model and exploration policy described in the main text.It frames these variations as potentially useful to researchers studying closely related models.Source: D. Model Variations Considered · p. 17

Transition Model

• The transition-model variations aimed to match consecutive latent scenes, but MONet’s changing slot order makes latent-space matching nontrivial and can misassign similar objects crossing paths.The authors preferred decoder-based pixel loss because it uses dynamic prediction simplicity to determine future slot ordering.Source: Transition Model · p. 17

Exploration Policy

• Rejection sampling was computationally expensive with sparse curiosity signals, while multidimensional Huber quantile regression trained slowly and scaled poorly with sparsity and action dimension.The deformation method trained quickly and worked well in artificial curiosity experiments even with 50-dimensional action spaces.Source: Exploration Policy · pp. 17–18

Curiosity Signal

• The authors used transition-model error as a rudimentary curiosity signal; stochastic-model variance was unsuitable because high variance in blank slots drowned out uncertainty about object motion.Blank-slot variance was amplified by their abundance, exceeding 50% of slots in the exploration environment; the authors therefore retained the simpler pixel-loss approach.Source: Curiosity Signal · p. 18

• Spriteworld contains occluding objects that vary in position, shape, and color; objects do not collide or interact otherwise, and their sizes are fixed.The environment is rendered as 64 × 64 × 3 images, and actions occupy the continuous hypercube [0, 1]^4.Source: E. Environment details · p. 18
• Actions select an object with a position click and move it with a motion click, adding Gaussian motion noise and clipping positions to keep objects in frame.The maximum motion per coordinate is 0.125, so an object can be moved between any two points within eight steps.Source: E. Environment details · p. 19
• Shape and position robustness tests change the target’s shape or place its initial position in a held-out lower-right quadrant, respectively.The position task includes a task-irrelevant distractor whose position remains unconstrained in both training and testing.Source: E. Environment details · p. 19
• Target-count and distractor-count robustness tests change the number of targets or distractors between training and testing, while retaining task reward over the targets.The target-count test changes from one target and two distractors to two targets and two distractors; the distractor-count test changes from one to two distractors with two targets.Source: E. Environment details · p. 19
• Sorting tests require objects of sampled colors to reach color-specific goals, with the red–blue pair held out; clustering tests hold out the purple–yellow pair after training on blue–green.The clustering reward uses the inverse Davies–Bouldin metric and gives a bonus for crossing its stated threshold.Source: E. Environment details · pp. 19–20

F. Sparse rewards

• Sparse rewards require deeper search rollouts or a value function; COBRA’s value-predictor version uses TD learning and selects actions with a one-step search maximizing V(zt+1).Dense shaped rewards had allowed a single-step search and reward predictor to suffice.Source: F. Sparse rewards · pp. 20–21
• With sparse terminal rewards, the value-predictor COBRA solves all tested situations, while the reward-predictor version reaches 50% success because smoothing gives near-goal scenes nonzero reward.Figure S6 shows the training curves for Goal Finding tasks with sparse terminal rewards.Source: Figure S6 · p. 21, F. Sparse rewards · p. 21

G. Vision Module Disentangling

• Six prominent latent dimensions yield disentangled object representations; the authors suggest this may help robustness to irrelevant-feature perturbations, though they did not explicitly encourage invariance.The labels on the latent traversals were assigned post hoc by eye.Source: G. Vision Module Disentangling · p. 21

H. Ablations and Additional Baselines

• Sweeping six prominent latent dimensions individually yields disentangled reconstructions; two additional non-coding “none” latents revert to their Gaussian prior.The latent labels were assigned post hoc by eye.Source: Figure S7 · p. 22
• MPO with either MONet features or the exploration policy falls between raw MPO and the combined MPO in data efficiency; exploration appears to give the larger boost to asymptotic performance and robustness.Figure S8 compares test-time performance and robustness as well as environment steps to reach 90% test performance.Source: Figure S8 · p. 22, H. Ablations and Additional Baselines · p. 22
• The task-time exploration ablation replaces COBRA’s policy with uniform actions in [0, 1]^4, while retaining exploration-policy use during the exploration phase.The data-efficiency plot for this comparison uses a linear rather than logarithmic scale.Source: Figure S9 · p. 23, H. Ablations and Additional Baselines · p. 23

I. Additional Transition Model Rollouts

• Figure S10 compares ground-truth observations with transition-model rollouts under the same action sequence, repeatedly moving the position click to the ground-truth object location.The vision-module decoder renders the model rollouts.Source: §I · p. 23, Figure S10 · p. 24