Object-Centric World Models from Few-Shot Annotations for Sample-Efficient Reinforcement Learning
Authors: Weipu Zhang
Adam Jelley
Trevor McInroe
Amos Storkey
Gang Wang
Venue: ICLR 2026 Poster (2026)
Keywords: Model-based RL
Object-centric RL
Video object segmentation
Atari
Hollow Knight
Primary area: reinforcement learning
TL;DR: We proposed an object-centric model-based RL pipeline, which integrates recent advances in computer vision to allow agents to focus on key decision-related elements. Extensive experiments demonstrate the efficacy of our method.

Abstract: While deep reinforcement learning (RL) from pixels has achieved remarkable success, its sample inefficiency remains a critical limitation for real-world applications. Model-based RL (MBRL) addresses this by learning a world model to generate simulated experience, but standard approaches that rely on pixel-level reconstruction losses often fail to capture small, task-critical objects in complex, dynamic scenes. We posit that an object-centric (OC) representation can direct model capacity toward semantically meaningful entities, improving dynamics prediction and sample efficiency. In this work, we introduce **OC-STORM**, an object-centric MBRL framework that enhances a learned world model with object representations extracted by a pretrained segmentation network. By conditioning on a minimal number of annotated frames, OC-STORM learns to track decision-relevant object dynamics and inter-object interactions without extensive labeling or access to privileged information. Empirical results demonstrate that OC-STORM significantly outperforms the STORM baseline on the Atari 100k benchmark and achieves state-of-the-art sample efficiency on challenging boss fights in the visually complex game **Hollow Knight**. Our findings underscore the potential of integrating OC priors into MBRL for complex visual domains.

Project page: https://oc-storm.weipuzhang.com

Published as a conference paper at ICLR 2026
OBJECT-CENTRIC WORLD MODELS FROM FEW-SHOT ANNOTATIONS
FOR SAMPLE-EFFICIENT REINFORCEMENT LEARNING
Weipu Zhang Adam Jelley!
Trevor McInroe' ?
Amos Storkey! Gang Wang
National Key Lab of Autonomous Intelligent Unmanned Systems, Beijing Institute of Technology
UUniversity of Edinburgh "Jiangxing Intelligence Inc.
weipuzhang academi e emaall com
{adam. jelley, t. t.m.iinnoee stw.ssyrrryye..... @ed ac.uk
gangwang@bit. adw.nngggiittttt.....mommmm edu. on
ABSTRACT
While deep reinforcement learning (RL) from pixels has achieved remarkable success, its sample inefficiency remains critical limitation for real-world ap- plications. Model-based RL (MBRL) addresses this by learning a world model generate simulated experience, but standard approaches that rely on pixel-level reconstruction losses often fail to capture small, task-critical objects in complex, dynamic scenes. We posit that object-centric (OC) representation rect model capacity toward semantically meaningful entities, improving dynam- ics prediction and sample efficiency. In this work, we introduce OC-STORM. object-centric MBRI. framework that enhances learned world model with object representations extracted by a pretrained segmentation network. By con- ditioning on minimal number of annotated frames, OC-STORM learns to track decision-relevant object dynamics and inter-object interactions without extensive labeling access to privileged information. Empirical results demonstrate that OC-STORM significantly outperforms the STORM baseline on the Atari 100k benchmark and achieves state-of-the-as sample efficiency on challenging boss fights in the visually complex game Hollow Knight. Our findings underscore the potential of integrating OC priors into MBRL for complex visual domains.
di-
Project page: https://oc-storne wwwpuzhang, com
INTRODUCTION
Deep reinforcement learning (RL) has achieved landmark successes in domains ranging from board games to robotic control (Silver et al., 2016 Mnih et al. 2015 Hafner et al. 2023; Feng et al. 2026). However, a fundamental limitation persists: Deep RL agents are notoriously sample-ineffciient, often requiring orders of magnitude more experience than humans to master a task. This inefficiency is particularly pronounced in pixel-based environments, where the agent must learn both visual representations and control policies from high-dimensional observations.
Model-based RL (MBRL) offers a promising path toward greater sample efficiency by learning predictive model of the environment dynamics (Sutton & Barto 2018 Ha & Schmidhuber 2018). Contemporary MBRL methods typically learn this world model self-supervised manner, using autoregressive prediction trained with pixel-wise reconstruction losses (e.g., the f2-loss) (Hafner al. 2023 Zhang et al. 2023 Micheli et al. 2023). While effective in many cases, this approach has critical weakness: the reconstruction objective is dominated by large, static background ele- ments, often the expense of neglecting small, sparse, yet decision-crities objects. As illustrated in Figure D standard world models like STORM can accurately reconstruct the background of visually complex game like Hollow Knight but fail to capture essential elements such as the boss character, which is key for successful policy learning, leading to poor performance.
Work was initiated at the University of Edinburgh and completed at the Beijing Institute of Technology. Coorrsponding author.

---

Published as a conference paper at ICLR 2026
Rewards, terminations
Environment
Observations
Pretrained video segmentation
World model
(a) A high-resolution sample observation from Hollow Knight.
Few-shot object anmotations
Learning imagination
Recompitatiin Gooandl Trath
Policy
Segreentatiio marks
Object features
Actions
(c) The proposed OC-STORM framework. A frozen, pre- trained segmentation model extracts object feature vectors from few annotated frames. These features are combined with downsampled pixels to train an OC world model, which is then used for policy learning via imagined trajectories.
(b) Reconstruction generated by trained STORM agent.
Figure 1: Left: STORM Zhang et al.) 2023) accurately reconstructs large background areas (blue) but overlooks the small, critical player and boss characters (orange), hindering policy learning. Right: Overview of the proposed OC-STORM framework. See Appendix A for network details.
A natural solution is to incorporate object-level inductive biases, guiding the world model to repre- sent the environment in terms of discrete, interacting entities. Historically, this required extensive task-specific labeling, limiting its practicality, Recent advances in open-set segmentation and track- ing models, including SAM (Krillov et al.) 2023 Ravi et al., (204)), Cutie (Cheng et al. (002)), and DIOO (Liu et al. 2023). have changed this landscape. These foundation models can generate high-quality object segmentations for novel domains with only handful of annotations, making OCRL viable and powerful approach.
In this work, we introduce OC-STORM, an object-centric MBRL framework that enhances world models with object representations extracted by a pretrained segmentation network (Figure Icp. Our approach begins by annotating key objects in handful of (e.g., 6 - 12) frames. A frozen, pretrained video segmentation network (e.g., Cutie (Cheng al. 2023)) then extracts compact feature vectors for these objects, which are fused with raw pixel observations to train a world model that explicitly reasons about object dynamics and inter-object interactions. The policy subsequently trained on imagined trajectories from the world model. By focusing model capacity on task-relevant entities, OC-STORM improves sample efficiency without requiring extensive labeling or access to internal environment states.
Our main contributions are summarized as follows.
c1) We propose OC-STORM, a novel MBRL framework which, to knowledge, is the first to successfully integrate few-shot, pretrained object segmentation models into world models for both the Atari 100k benchmark and the visually complex game Hollow Knight. This general applicability is achieved without extensive labeling or access to in- ternal state information.
c2) We provide comprehensive empirical evaluation across diverse domains (Atari, Hol- low Knight), model backbones (STORM, DreamerV3), and segmentation methods (Cutie, SAM2). Our results show that OC-STORM achieves state-of-the-art sample efficiency, particularly in environments where key information localized in objects.
c3) We conduct extensive ablation study that yields practical insights for future OCRL meth- ods, including comparisons between vector- rector-based and mask-based object representations and an analysis of robustness segmentation errors.
2

---

Published as a conference paper at ICLR 2026
2 PRELIMINARIES
We formulate the game control problem as a finite-horizon Markov decision process (MDP) M =
(8, A. P. T, 7), with latent state space S, action space A, transition kernel p(*-11 St: ae), reward function r S x A + R, and discount factor y E [0, 1). Our objective is to learn a policy To(a; (tt) parameterized by 8 that maximizes the expected return Ex E 1- y'r:]. We adopt standard MBRL paradigm: first, learn a world model Pe via self-supervised learning on environment interactions for, 0t, rz); then, train the policy To using trajectories generated by the model.
To extract object information from out-of- domain environments (e.g., Atari) with minimal human effort, evaluated several otate-of-the-art the-art object segmentation, detection, and tracking methods
CCeng & Schwwnng 2022 Cheng et al. 2023 Kirillov et al. 2023; /Zhang et al.
2024
Ravi et al.
2024 Redmon al. 2016 Jocher al. 2023 Liu al. [0023] Locatello et al.
2020
Kipf al.
2022
Elsayed et al.) 2022 [ Wang et al. 2203 |
Xu et al. 2023); see Appendix for a detailed review.
Among these, we selected Cutie (Cheng al., 2023) and SAM2 (Ravi et al. 2024) as our object
feature extractors due to their suitability for our setting. These methods exhibit four key properties
that are critical for our task:
Temporal consistency and efficiency. As video-based segmentation methods, they pro- duce consistent object representations across frames and support real-time processing of high-resolution inputs on a single consumer-grade GPU.
Retrieval- flexibility. They effectively utilize few-shot annotations through retrieval from a memory bank, which crucial complex environments where a single frame or natural language description may not capture all object states.
Compact object representations. They provide compact vector embeddings that preserve sufficient semantic and visual information necessary for decision decision-making.
Out-of-domain robustness. Despite being trained on natural images without exposure to Hollow Knight or Atari frames, they generalize effectively to Atari and Hollow Knight frames, as demonstrated in our experiments.
Both SAM2 and Cutie employ an object-attention mechanism that facilitates object-level and visual-level information interaction. This design leverages high-level semantic information to im- prove segmentation performance. We harness this mechanism to extract compact object-leve feature representations for our world model.
3
METHOD
Our method consists of two stages: (1) self-supervised learning of an OC world model that captures environment dynamics, and (2) training an actor-critic policy using trajectories imagined by the model. We denote the world model parameters as Φ, the critic network parameters as v. and the actor network parameters as We use L to represent the batch length of trajectory segments and T for the episode length. The model architecture is detailed in Appendix
3.1 OBJECT FEATURE AND VISUAL INPUT
We extract object feature representations using the object-attention outputs from SAM2 or Cutie. The dimension of these object features is denoted obj dim (256 for SAM2 and 048 for Cutie). For visual input, we resize the original observation to 64 x 64 resolution, consistent with prior work (Hafner al. 2023
Zhang et al., 2023). Formally, at each timestep 4, we define:
Observation:
Ot E x3xHxW
Object features:
St obj SegModel(o:) E Koobj.ddi
(1)
Visual input:
S Resize(o.) 33x64X64
where K is the number of objects, which is enpiron and set by the user (e.g., K 3 for Pong: two paddles and one ball). The segmentation model (SAM2 or Cutie) maintains internal states across frames to ensure tracking consistency, which reset the start of each episode.
3

---

Published as a conference paper at ICLR 2026
3.2 LATENT STATE DISCRETIZATION VIA CATEGORICAL VAE
Autoregressive sequence models operating directly on high-dimensional inputs are susceptible to compounding prediction errors (Hafner al. 2023 Zhang al. 2023). To mitigate this, we employ a categorical VAE (Kingma & Welling) 2014] Hafner et al. 2023) that encodes inputs into discrete
latent representations. The encoder 9e maps input states s, to discrete latent variables 2t, while the decoder pa reconstructs the inputs from these latents:
Encoder: zt ~ (0(2))));
(2)
Decoder: $2 =o(z)
We use distinct architectures for different input modalities: multi-layer perceptron (MLP) en- coders/decoders for object feature vectors (IR Xdaas ++ 1g x1xx161 and convolutional neural net- IR 3232232). The object latents use 16 categorical distributions with 16 classes each, while visual latents use 32 distributions with 32 classes, following established configurations (Hafner et al., 2023) Zhang et al. 2023). The decoder architectures are symmetric to the encoders. The reduced dimensionality for objects reflects their lower information content compared to full scenes. We use the straight-through gradient estimator Bengio et al. 2013) to enable differentiable sampling from the categorical distributions.
work (CNN)-based ones for visual observations (R3x64x64
3.3
SPATIAL- TEMPORAL OBJECT-CENTRES DYNAMICS
Our framework supports both transformer-based (STORM) and recurrent neural network (RNN)- based (DreamerV3) world model backbones. The core innovation is that separately models object and visual dynamics while enabling interaction between them.
architecture
For the STORM backbone, we employ a transformer architecture with alternating spatial and tem- poral attention blocks. Spatial attention operates across the object tokens and one visual token at each timestep t: (z/ 2?:........................
..... 2x 2xie), capturing inter-object relationships and object-scene in-
teractions. Temporal attention then processes each token type independently across the sequence (=t =2, ... = to model dynamics. Actions are incorporated via concatenation with the latent rep- resentations. For DreamerV3, we augment the RNN with spatial attention mechanisms at each timestep achieve similar object-scene interaction.
The sequence model fe takes as input the sequence of object latents 21.1. visual latents 2 and actions 01:L, and outputs hidden states h1.L E () as
h (= ohj 21:L: 01.1) )...)) E R(K+1)xLxd,
(3)
where d2 = 256 denotes the hidden dimension, and K -1 accounts for the K object tokens plus one visual token.
3.4 PREDICTION HEADS
We use the hidden states fú2 of the transformer to predict the next latent state, reward, and termination signal. The dynamics predictor 96 Dyy is an MI.P that predicts the distribution of the next latent state 2441- For reward and termination prediction, we use a self-attention mechanism that aggregates information from all objects and visual tokens, with structures detailed in Appendix A Specifically, introduce a special query token (similar to the [CLS] token in BERT (Devlin et al., (009)) that
attends to all tokens in h, and then passes the resulting representation through an MLP to predict the reward fü and termination probability î. These predictors are given as follows:
Dynamics predictor:
9 dyn
d (2c4 (
Reward predictor:
f. = rew (ht);
(4)
Termination predictor:
Ar = g term (1+).
3.5 TRAINING THE WORLD MODEL AND POLICY
We train the world model end-to-end using combination of reconstruction and prediction losses. The world model is trained to maximize the likelihood of the observed data, including the recon- structed inputs, rewards, and terminations. The policy my and value function Vp trained using
4

---

Published as a conference paper at ICLR 2026
imagined trajectories generated by the world model. We use the actor-critic algorithm from Dream-
erV3 (Hafner et al. 2023 Zhang et al. 2023 Micheli et al. 2023; Robine et al., (002)), which
involves policy gradient updates with value function baselines. Detailed loss functions and training procedures are provided Appendix
4 EXPERIMENTS
We evaluate OC-STORM two challenging benchmarks: the standard Atari 100k sample- efficiency benchmark (Bellemare et al., 2013) and complex boss fights from Hollow Knight (Team- Cherry 2017). highly acclaimed game released in 2017. This dual evaluation tests method across diverse visual complexities-fromt the relatively simple but well-established Atari environ- ments to the rich, dynamic visuals of modern games where OC reasoning provides maximum benefit.
scores the Atari 100k benchmark. We compare OC approaches (vector-bas OC and mask-based FOCUS (Ferraro et al. (023)) against baseline world models. Bold indicates scores within 5% of best. The highlighting separately with respect to the Dreamer and STORM baselines, shown in blue and red, respectively. The last column '#obj' denotes manually annotated objects per game. All experiments use identical lightweight configurations suitable for nigh-resolution environments. STORM serves as the primary baseline due its computational efficiency.
Table 1: Game scores and overall
computed
Vector
Mask FOCUS STORM
Voctor SAM2-OC
Vector
Cutie-OC DreamerV3
Cutie-OC STORM
Game
Random
Human
DreamerV3
STORM
STORM
#obj
Alien
228
7128
789
1234
748
1039
591
1101
Amidar
1720
107
123
110
150
163
Assault
222
742
928
1201
1377
986
1086
1270
Asterix
210
8503
916
1293
1318
1552
1408
1754
3
HankHeist
753
812
849
990
702
840
1075
3
BattleZone
2360
37188
7500
9550
5830
12870
12950
4590
Boxing
0
12
78
88
81
78
81
92
Breakout
39
57
52
3
ChopperCommant
811
7388
1926
2130
1644
1751
2580
2090
CrazyClimher
10780
35829
83635
89671
79196
65753
76039
84111
DemonAttack
152
1971
379
540
325
360
491
411
Freeway
0
Froshite
65
4335
272
275
3666
277
294
260
Gopher
258
2412
4752
4244
5307
4453
2690
4457
2
Hero
1027
30826
6777
3712
11434
9632
7913
6441
Jamesbond
303
544
470
408
418
383
347
Kangaroo
52
3035
1322
3300
3512
2418
3141
4218
1598
2666
7111
7060
6522
6927
10410
9715
KungFuMaster
258
22736
11055
22863
20046
20489
23036
24988
MsPacman
307
6952
1720
1822
1490
1557
2464
2401
2
Pong
-21
18
14
PrivateEye
69571
216
364
Qbert
164
13455
3652
2974
2910
3220
4256
4546
RoadRunner
7845
15075
18014
14841
12841
16853
20482
Seaquest
42055
653
712
UpNDown
533
11693
5009
7486
6128
5098
47
6623
HNS mean
0%
100%
106.6 666.6%
119.4%
114,2%
107.2%
124.6%
134.8%
HNS median
0%%
100%
33.2%
42.6%
42.5 %
35.5%
35.0%
43.8%
Obj-detectable
144.0%
160.8%
147.7%
142.4%
175.9%
186.2%
Otherwise
69.2%
78.0%
80.8%
72.0%
73.3%
83.4%
4.1 ATARI L00K
We evaluate our method under the standard Atari 100k protocol (Kiserr et al. 2020). which limits training to 100, 000 environment frames. For each game, we report average performance across 5 random seeds, with each seed evaluated over 20 episodes. Performance measured using human- normalized scores (HNS): HNS (agent random) (human random). To isolate the impact of OC representations, we focus here on ablations within our framework. External benchmarks against other MBRL approaches are presented in Appendix
5

---

Published as a conference paper at ICLR 2026
We conduct experiments with two typical world modeling backbones: DreamerV3 (Hafner al.
2023 and STORM (Zhang et al. 2023); two object extraction vision models: SAM2 (Hafner et al.
2023 and Cutie (Cheng et al. 2023); and two object representation methods: vector-based and mask- based. The mask-based representation follows the idea proposed in FOCUS (Ferraro et al, 2023); in our implementation, it is realized by feeding Cutie's segmentation masks into STORM. The results are shown in Table [ Overall, OC methods clearly outperform the baseline. Cutie- based variants achieve higher scores than SAM2-based ones, while the mask-based representations underperform vector-based ones and are only on par with baselines. We therefore select the best performing variant, Cutie-OC-STOMM, in terms of HNS mean and median as our proposed algorithm, OC-STORM.
To further examine the effectiveness of our method, we categorize the 26 games into two groups: those where all decision-relevant objects can be consistently identified by SAM2 or Cutie, and those where detection is incomplete | For the second group, we further discuss the limitations of current vision models in Section 7 The results, reported at the bottom of Table I show that vector-based OC methods substantially outperform the baseline in the first category, while in the second category they still achieve performance comparable the baseline. These findings highlight both the benefits of OC learning and the robustness our framework.
Interpretation of the results Although SAM2 achieves stronger results than Cutie video seg- mentation benchmarks, our OC agents benefit more from Cutie's ' features. This discrepancy arises from their object-feature design: Cutie aggregates visual features within masked regions, whereas SAM2 produces a prototype vector intended for classification. Consequently, SAM2 features may lose positional and global context, providing weaker guidance for policy learning.
FOCUS-like mask-hased methods mainly suffer from resolution issues: low-resolution inputs (e.g., 64 64 in world models) discard crucial object details, while high-resolution masks incur quadratic memory and computation costs. Moreover, raw masks are noisy, and thus their dynamics are difficult to predict. In contrast, the vector representation semantically summarized from high-resolution input, which is more consistent and computationally efficient. Additional training curves comparing different representations are provided in Figure 4a
4.2 HOLLOW KNIGHI
While the Atari benchmark is widely used in the RL community, is not sufficient verify the bustness of our proposed method. Atari's visual simplicity, where object pixels can often be sepa- rated from the background by well-defined pixel- value boundaries, allows methods like Dream- erV3 and STORM to capture the environment dy- namies almost perfectly in raw visual space. Such simplicity would be rarely seen in real-world sce- narios. In contrast, the boss fights in Hollow Knight offer more suitable testbed, where the visuals are much more complex, including many dynamic and distracting elements
GodTamer
DommetPrutectoe
MageLord
40
30
20
20
50k
100k
5ik
100k
50kk
10%k
MantisLord
Mawlek
Primeime
20
20
50k
100k
Sik
L00k
Sikk
100k
STORM
OC-STORM
For Hollow Knight, we similarly limit the number of samples to 100k, equivalent to approximately 3.1 hours of real-time gameplay at 9Hz. For each boss, we conduct three experiments with different random seeds.
Figure 2: The training episode returns on Hol- low Knight. We use solid line to represent the mean of 3 seeds and use semi-transparent background to represent the standard deviation.
'The first category "Objideetectabee)) includes Asterix, BankHeist, BattleZone, Boxing, ChopperCom- mand, DemonAttack. Jamesbond, Kangaroo, Krull, Pong. RoadRunner, Seaquest, and UpNDown. The sec- ond category includes Alien, Amidar, Assault, Breakout, CrazyClimher, Freeway, Frosthite, Gopher, Hero, KungFuMaster, MsPacman, PrivateFye, and Qbert. Please refer the evaluation videos in the supplementary materials.

---

Published as a conference paper at ICLR 2026
As seen in Figure though the original STORM also learn good policy, proposed OC method converges significantly faster and yields stronger performance most cases, especially when the environment is more challenging, such for Mage Lord and Pure Vessel. Numerical results are shown in Appendix
F.1
Since Hollow Knight is yet established benchmark, existing methods differ significantly in sample step limits, environment wrapping, reward functions, boss selection, among others. This makes direct comparisons with existing methods challenging. As the primary goal of this work is improve MBRL through the use of OC representations, we therefore compare our results with the equivalent baseline algorithm STORM. Nevertheless, we include the results from Yang (2023) on the boss Hornet Protector for rough comparison in Appendix E.6
ANALYSIS
5.1 COMPLETENESS OF THE OBJECT REPRESENTATION
We utilize the output feature of the object transformer in the video model. While this feature theo- retically contains all the state and positional information of an object, it is uncertain whether fully captures these details in practice. Specifically, need to determine if the masked pooling could potentially obscure positional information. The agent's 's performance, as demonstrated in Section 4 provides general quantitative evidence. Here, present qualitative evidence support this claim.
We trained a 4-layer Conv TonvTraaspose2 (Zeiler et al. 2010) decoder on the Atari Boxing game. takes vector-based object features as inputs, corresponding to the white and black players, respectively, to reconstruct the observation. The dataset was collected using a random policy, with 10, 000 frames for training and 000 frames for validation. Sample reconstructions result from the validation set are shown in Figure Ba This indicates that these features effectively capture the state and position of the objects. Both SAM2 and Cutie features can provide similar outcomes in this test, but the model takes a larger number of update steps to converge when using the SAM2 feature compared using the Cutie feature.
Boxing
Ground Truth
100
50k
100k
Object Mask
Pong
20-
0-
-20
Prodiction
5ök
100k
0%
25%
12.5%
50%
(a) Reconstructions with object features.
(b) Robustness analysis.
Figure Analysis of the object feature. (a) Observation reconstructions on Atari Boxing with two object feature vectors as inputs. The object mask row highlights the relevant objects. (see Section 5.1 for details). (b) Training episode returns for Atari Boxing and Pong with different zeroing probabilities (see Section 5.2)
5.2 ANALYSIS OF SEGMENTATION MODEL FAILURES
To mimic segmentation model failures, we randomly set the object feature vector to 0. This opera- tion is identical to how process the feature when the video model detects nothing, described in

---

Published as a conference paper at ICLR 2026
Appendix G.2 We conduct experiments Atari Boxing and Pong with different zeroing probabil- ities. To avoid interference from visual imputs, we only use the object module in these experiments.
The results presented in Figure 3b As the detection accuracy of the vision model increases, the agent's performance improves accordingly. This also demonstrates the robustness OC-STORM in handling unstable detection results. Additionally, since the zeroing process is purely random and the agent is trained only after the termination of each episode, every new episode during training serves an indicator of test-time failure performance.
5.3 ADDITIONAL EXPERIMENTS WITH CONTINUOUS CONTROL
To evaluate the potential of OC-STORM on continuous control tasks, we conduct four experiments on the Meta-world (Yu et al. 2019) benchmark. We compare our results with MWM (Seo et al. 2022), which is also designed to help the world model focus on small dynamic objects. We choose one easy, two medium, and one hard tasks according to the MWM paper (see Seo et al. (2022) Appendix F. Experiments Details). These tasks are randomly selected to cover different objects and policies.
As shown in Figure 4b OC-STORM generally exhibits higher sample efficiency than STORM. In some tasks, also outperforms MWM in terms of efficiency and performance. This provides evidence that this approach can also perform well continuous tasks out-of-the-box, without sig- nificant adaptation of the pipeline or extensive tuning for these very different continuous control environments.
Boxing
Pong
reach
sweep
100
100%
100%
75%
75%
50
0-
50%
50%
25%
25%
0
-20
0%
0%
50k
100k
0
50k
100k
0
100k
200k
200k
400k
MageLord
HornetProtector
basketball
pick-out-of-hole
100%
100%
30
75%
75%
50%
50%
25%
50k
100k
50k
100k
200k
400k
500k
1000k
OC STORM
STORM
OC-STI STORM
MWM
FOCUS
No-visual
STORM
DreamerV2
(a) Module ablation training curves.
(b) Metaworld training curves.
Figure 4: (a) Training episode returns under different input module configurations. STORM utilizes only the visual module, "No-visual" relies solely on the object module, and OC-STORM incorpo- rates both, as described in Appendix A FOCUS employs mask-based representations. The top two tasks are from Atari, while the bottom two are from Hollow Knight. See numerical results in Sec- tion 4.1 (b) Training success rates on Meta-world tasks. The results of MWM (Seo et al. 2022)
and DreamerV2 (Hafner et al. (021) are from the MWM paper. See details in Section 5.3
5.4 EXTRA ANALYSIS
We provide additional ablations in Appendix F examining the policy design choices, the number of annotations, and the computational overhead of introducing pretrained vision models.
RELATED WORK
Model-based RL Ha & Schmidluuer (2018) first demonstrated the feasibility of learning by imag- ination in pixel-basee based environments. SimPLe (Kiserr et al.) al, (2000) further extended this idea to Atari games (Bellemare et al. 2013), though with limited efficiency. The Dreamer series (Hafner al.,
8

---

Published as a conference paper at ICLR 2026
2019 2021: 2023) employs categorical VAEs and RNNs, to achieve robust performance across di- verse domains. Dreamer introduces both a stable discretization method and a set of techniques for robust optimization across domains with diverse observations, dynamics, rewards, and goals. TWM (Robine et al. 2023) and STORM (Zhang et al. 2023) replace the RNN sequence model in Dreamer with transformers, enhancing parallelism during training. TWM encodes the observation, reward, and termination as three input tokens for the transformer, while STORM encodes them as a single token, demonstrating better efficiency. IRIS (Micheli et al. (2023), and improved efficiency vari- ants A-IRIS (Micheli et al. (024) and REM (Cohen et al., 2024), utilize VQ-VAE (van den Oord al., 2017) for multi-token latent representations. DIAMOND (Alonso et al., 2024) employs a diffusion process the world model, further improving the final performance. All of these methods predominantly use an L2 reconstruction loss for self-supervised learning.
Object-centric RL Many attempts have been made introduce OC learning to RL systems. How- ever, to our knowledge, no existing methods have been directly applied to Atari games Hollow Knight without leveraging internal game states or an extensive number of annotations. These OC learning methods broadly follow two main trends: two-stage and end-to-end.
Two-stage methods usually first employ computer vision techniques to detect objects, then train the policy based on this object-level information. Current approaches often require labour-heavy task-specific fine-tuning (Devin et al., 2018; Liu al. 2021), access to game memories (Delfosse) cet al. 2023; Jain 2024). or leverage game-specific observation structures (Stanic et al. (024). FOCUS (Ferraro et al., 2023), > the most similar work to ours, is model-based method that uses Trathing ( Yang et al. 2023) to generate segmentation masks, which are then fed into DreamerV2 (Hafner et al. 2021) for policy training. Using binary masks for object representa- tion limits efficiency, which was discussed in Section [ Moreover, FOCUS has only been tested on six ttean-backgroud fixed-camera roboties manipulation tasks, and hasn't 't been fully explored in more visually complex environments.
End-to-end methods jointly learn the object perception and policy, often using unsupervised slot- based approaches (Locatello et al., 2020) discover and represent objects. While these methods allow the visual module to be trained alongside the world model or policy network, their unsuper- vised learning nature leads to poor object detection quality, especially noisy, real-worke scenes. As
result, they are typically limited to visually simple OC benchmarks (Watters et al. [024]] Ahmed
al., (02T) and struggle to be adapted to visually complex tasks. Several model-based (Veerapa- neni et al.) 2019 Lin et al.) 22000] van Bergen & Lanillos, 2022) and model-free (Yoon et al. [0023]
Haramati et al. 2024) algorithms have used these ideas. Nakano et al. (2024) added slot attention to STORM, achieving stronger performance the OCRL benchmark (Yoon et al. 2023
LIMITATIONS
Our method has two main limitations, each of which corresponds to a potential future direction.
Duplicated instances Current video object segmentation algorithms primarily developed and trained to track a single object. When a scene contains two or more identical objects, approaches like SAM2 or Cutie may fail to segment each object correctly and thus may affect performance.
Geometric map representation Our object representations are not well-suited for encoding geo- metric structures such walls, boundaries, and navigable spaces. This limitation necessitates taining raw visual inputs in our pipeline and reflects a broader challenge for current OC approaches in the visual domain.
These limitations are further illustrated and explained in Appendix
8 CONCLUSIONS
In this work, we introduced OC-STORM, an MBRL pipeline designed to improve sample efficiency in visually complex environments, By integrating recent advances in object segmentation and detec- tion, we mitigate the limitations of traditional reconstruction-based MBRL. methods, which may be

---

Published as a conference paper at ICLR 2026
dominated by large background areas and overlook decision-relevant details. Through experiments across various domains and model configurations, we demonstrated that OC learning could be suc- cessfully implemented without relying internal game states extensive labelling, highlighting the adaptability of our method to complex, visually rich environments. OC-STORM represents a meaningful step toward combining modern computer vision with RI, offering efficient frame- work for training agents in visually complex settings.
ACKNOWLEDDGEMENTS
The work of W.Z. and G. G.W. in this paper was supported by the National Natural Science Foundation of China under Grant U23B2059. The work of A.S. part of project of the European Union's 's Hori- zon Europe research and innovation programme under Grant Agreement No. 101120726, funded by UK Research and Innovation (UKRI) under the UK govermment's Horizon Europe funding guaran- tee 10085198.
REFERENCES
Ossama Ahmed, Frederik Träuble, Anirudh Goyal, Alexander Neitz, Manuel Wuthrich, Yoshua Ben- gio, Bernhard Schölkopf, and Stefan Bauer. CausalWorld: A Robotic Manipulation Benchmark for Causal Structure and Transfer Learning. In International Conference on Learning Represen- tations, 2021.
Eloi Alonso, Adam Jelley, Vincent Micheli, Anssi Kanervisto, Amos Storkey, Tim Pearce, and François Flcuret. Diffusion for World Modeling: Visual Details Matter Atari. CoRR, abs/2405.12399, 2024.
M. Bellemare, Naddaf, J. Veness, and M. Bowling, The Arcade Learning Environment: An Evaluation Platform for General Agents. Journal of Artificial Intelligence Research, 47:253-279, June 2013. ISSN 1076-9757.
Yoshua Bengio, Nicholas Léonard, and Aaron C. Courville. Estimating or Propagating Gradients Through Stochastic Neurons for Conditional Computation. CoRR, abs/1308.3432, 2013.
Yusuf Bham and Jonathon Wyza, hk-modding/api, August 2017. URL httttt://tthhuubb com/
hk-modd r /
Greg Brockman, Vicki Cheung. Ludwig Pettersson, Jonas Schneider, John Schulman, Jie Tang, and Wojciech Zaremba. OpenAl Gym. CoRR, Abss/600.01540 2016.
Ho Kei Cheng and Alexander G. Schwing. XMem: Long-term Video Object Segmentation with an Atkinson-shiffrin Memory Model. In European Conference on Computer Vision, volume 13688 of Lecture Notes in Computer Science, pp. 640-668. Springer, 2022.
Ho Kei Cheng, Seoung Wug Oh, Brian L. Price, Joon-Young Lee, and Alexander G. Schwing. Putting the Object Back into Video Object Segmentation. CoRR, 1tt//230..12822. 2023.
Lior Cohen, Kaixin Wang, Bingyi Kang, and Shie Mannor. Improving Token-based World Models with Parallel Observation Prediction. CoRR, abs/2402.05643 2024.
Sarah Cooley, Introduction to Hyper-V V on Windows 10, April 2022. URL https: www.wwwwwwwwww....w.......ww.www /abutt/
CrankyTemplar. Hollow Knight Sisters of Battle (Radiant Difficulty, Nail Only, No Damage),
August 2018. URL https: youtube com/ maatch?v2200XXxxxx
Ruifeng Cui. allco623/DDONNhollvvmiii
2021.
URL https://github. com/
A1lec0623/DON Hollow HollowKnight
Quentin Delfosse, Jannis Blüml, Bjarne Gregori, Sebastian Sztwiertnia, and Kristian Kerst-
ing. OCAtari: Object-centric Atari 2600 Reinforcement Learning Environments. CoRR,
abs/2306.08644 2023.
10

---

Published as a conference paper at ICLR 2026
Coline Devin, Pieter Abbeel, Trevor Darrell, and Sengey Levine. Deep Object-centric Represen- tations for Generalizable Robot Learning. In IEEE International Conference on Robotics and Automation, pp. 7111-7118. IEEE, 2018.
Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding. In Proceedings of Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, pp. 4171-4186 Association for Computational Linguistics, 2019.
Stefan Elfwing, Eiji Uchibe, and Kenji Doya. Nigmoid-weightee linear units for neural network function approximation in reinforcement learning. Neural Networks, 107:3-11, 2018.
Gamaleldin F. Elsayed, Aravindh Mahendran, Sjoerd van Steenkiste, Klaus Greff, Michael C. Mozer, and Thomas Kipf. SAVi++: Towards End-to-end Object-centric Learning from Real- world Videos. In Advances in Neural Information Processing Systems, 2022.
Zhaohan Feng, Ruiqi Xue, Lei Yuan, Yang Yu, Ning Ding, Meiqin Liu, Bingzhao Gao, Jian Sun, Xinhu Zheng, and Gang Wang. Multi-agent embodied AI: Advances and future directions. Sci- ence China Information Sciences, 69(5):151202, 2026.
Stefano Ferraro, Pietro Mazzaglia, Tim Verbelen, and Bart Dhoedt. FOCUS: Object-Centric World Models for Robotic Manipulation. In aitinnsicall-Mooivvreee and Open- Open-Ended Learning Work- shop @ NeurIPS2023, November 2023.
David Ha and Jürgen Schmidhuber. World Models. CoRR, 0bs//1033.1122 2018.
Danijar Hafner, Timothy P. Lillicrap, Ian Fischer, Ruben Villegas, David Ha, Honglak Lee, and James Davidson. Learning Latent Dynamics for Planning from Pixels. In Proceedings of In- ternational Conference on Machine Learning, volume 97 of Proceedings of Machine Learning Research. pp. 2555-2565, PMLR, 2019.
Danijar Hafner, Timothy P. Lillicrap, Mohammad Norouzi, and Jimmy Ba. Mastering Atari with Discrete World Models. In International Conference on Learning Representations, 2021.
Danijar Hafner, Jurgis Pasukonis, Jimmy Ba, and Timothy P. Lillicrap. Mastering Diverse Domains through World Models. CoRR, has/2301.041044 2023.
Dan Haramati, Tal Daniel, and Aviv Tamar. Entity-centric Reinforcement Learning for Object Ma- nipulation from Pixcls. In International Conference on Learning Representations, 2024.
Matteo Hessel, Joseph Modayil, Hado van Hasselt, Tom Schaul, Georg Ostrovski, Will Dabney, Dan Horgan, Bilal Piot, Mohammad Gheshlaghi Azar, and David Silver. Rainbow: Combining Improvements in Deep Reinforcement Learning. In Proceedings of AAAI Conference on Artificial Intelligence, pp.
AAAI Press, 2018.
Aditya Jain. AdityaJain 1330/HKRI, June 2024. AdityaJain 0000/1K10001111
URL. https: com/
Glenn Jocher, Ayush Chaurasia, and Jing Qiu. Ultralytics YOLOv8, 2023.
URL https:
//github. www./ooot/alooooo.ooooo..oolaala htt.b//..cco//lllraaartt///////s wtw.://ralytcccs//ttaal.ss.. lyiiss
Lukasz Kaiser, Mohammad Babaeizadeh, Piotr Milos, Blazej Osinski, Roy H. Campbell, Konrad Czechowski, Dumitru Erhan, Chelsea Finn, Piotr Kozakowski, Sergey Levine, Afroz Mohiuddin, Ryan Sepassi, George Tucker, and Henryk Michalewski. Model Based Reinforcement Learning for Atari. In 8th International Conference on Learning Representations, ICLR 2020, 2020.
Diederik Kingma and Jimmy Ba. Adam: A Method for Stochastic Optimization. In International Conference on Learning Representations, 2015.
Diederik P. Kingma and Max Welling. Auto-encoding Variational Bayes. In International Confer- ence on Learning Representations, 2014.
Thomas Kipf, Gamaleldin Fathy Elsayed, Aravindh Mahendran, Austin Stone, Sara Sabour, Georg Heigold, Rico Jonschkowski. Alexey Dosovitskiy, and Klaus Greff Conditional Object-centric Learning from Video. In International Conference Learning Representations, 2022.

---

Published as a conference paper at ICLR 2026
Alexander Kirillov, Eric Mintun, Nikhila Ravi, Hanzi Mao, Chloé Rolland, Laura Gustafson, Tete Xiao, Spencer Whitchead, Alexander C. Berg, Wan- Yen Lo, Piotr Dollár, and Ross B. Girshick. Segment Anything. In /EEE/CVF International Conference Computer Vision, pp. 3992-4003. IEEE, 2023.
J. (Jennifer) Lee. Adapting Model-Free Reinforcement Learning for Boss Fights in Hollow Knight: reward shaping approach. Bachelor's 's thesis, LIACS, Leiden University, 2023.
Zhixuan Lin, Yi-Fu Wu, Skand Vishwanath Peri, Bofeng Fu, Jindong Jiang, and Sungjin Ahn. Im- proving Generative Imagination in Object-centric World Models. In Proceedings of International Conference on Machine Learning, volume 119 of Proceedings of Machine Learning Research, pp. 2140-6149, PMLR, 2020.
Iou-Jen Liu, Zhongzheng Ren, Raymond A. Yeh, and Alexander G. Schwing. Semantic Tracklets: An Object-centric Representation for Visual Multi-agent Reinforcement Learning. In /EEE/RSJ International Conference on Intelligent Robots and Systems, pp. 5603-5610. IEEE, 2021.
Shilong Liu, Zhaoyang Zeng. Tianhe Ren, Feng Li, Hao Zhang. Jie Yang, Chunyuan Li, Jianwei Yang. Hang Su, Jun Zhu, and Lei Zhang. Grounding DINO: Marrying DINO with Grounded Pre-training for Open-set Object Detection. CoRR, abs/2303.05599, 2023.
Francesco Locatello, Dirk Weissenborn, Thomas Unterthiner, Aravindh Mahendran, Georg Heigold, Jakob Uszkoreit, Alexey Dosovitskiy, and Thomas Kipf. Object-centric Learning with Slot At- tention. In Advances in Neural Information Processing Systems, 2020.
James Lucas, George Tucker, Roger B. Grosse, and Mohammad Norouzi. Understanding Posterior Collapse in Generative Latent Variable Models. Deep Generative Models for Highly Structured Data Workshop, 2019.
Vincent Micheli, Eloi Alonso, and François Fleuret. Transformers are Sample-efficient World Mod- cls. In International Conference on Learning Representations, 2023.
Vincent Micheli, Eloi Alonso, and François Fleuret. Efficient World Models with Context-aware Tokenization. CoRR, abs/2406. 19322006....0..... 19320, 2024.
Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Andrei A. Rusu, Joel Veness, Marc G. Belle- mare, Alex Graves, Martin A. Riedmiller, Andreas Fidjeland, Georg Ostrovski, Stig Petersen, Charles Beattie, Amir Sadik, Ioannis Antonoglou, Helen King, Dharshan Kumaran, Daan Wier- stra, Shane Legg, and Demis Hassabis. Human-level control through deep reinforcement learning. Nature, 518(7540):529 533, 2015.
Philipp Moritz, Robert Nishihara, Stephanie Wang. Alexey Tumanov, Richard Liaw, Eric Liang, Melih Elibol, Zongheng Yang, William Paul, Michael I. Jordan, and Ion Stoica. Ray: A Dis- tributed Framework for Emerging AI Applications. In USENIX Symposium on Operating Systems Design and Implementation, pp. 561-577. USENIX Association, 2018
Akihiro Nakano, Masahiro Suzuki, and Yutaka Matsuo. Learning Compositional Latents and Behav- iors from Object- I Latent Imagination. In The Annual Conference of the Japanese Society for Artificial Intelligence, 2024.
Nikhila Ravi, Valentin Gabeur, Yuan- Main-Ting&n Hu, Ronghang Hu, Chaitanya Ryali, Tengyu Ma, Haitham Khedr, Roman Rädle, Chloe Rolland, Laura Gustafson, Eric Mintun, Junting Pan, Kalyan Va- sudev Alwala, Nicolas Carion, Chao- Yuan Wu, Ross Girshick, Piotr Dollár, and Christoph Fe- ichtenhofer. SAM 2: Segment Anything in Images and Videos, August 2024. URL http:
7/ aaaaiivvv..om wrw. abs/2408 2000/1100000001444.0 00714
Joseph Redmon, Santosh Kumar Divvala. Ross B. Girshick, and Ali Farhadi. You Only Look Once: Unified, Real-time Object Detection. In IEEE Conference on Computer Vision and Pattern Recognition, pp. 779-788. IEEE Computer Society, 2016.
Jan Robine, Mare Höftmann, Tobias Uelwer, and Stefan Harmeling. Transformer- baseed World Mod- els Are Happy With 100k Interactions. In International Conference on Learning Representations, 2023.
12

---

Published as a conference paper at ICLR 2026
John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal Policy 2017.
Optimization Algorithms. CoRR,
Max Schwarzer, Ankesh Anand, Rishab Goel, R. Devon Hjelm, Aaron C. Courville, and Philip Bachman. Data-efficient Reinforcement Learning with Self-predictive Representations. In 9rh International Conference Learning Representations, 2021.
Younggyo Seo, Danijar Hafner, Hao Liu, Fangchen Liu, Stephen James, Kimin Lee, and Pieter Abbeel. Masked World Models for Visual Control. In Conference on Robot Learning, volume 205 of Proceedings of Machine Leurning Research, pp. 1332-1344. PMLR, 2022.
David Silver, Aja Huang, Chris Maddison, Arthur Guez, Laurent Sifre, George van den Driess- che, Julian Schrittwieser, loannis Antonoglou, Vedavyas Panneershelvam, Marc Lanctot, Sander Dieleman, Dominik Grewe, John Nham, Nal Kalchbrenner, Ilya Sutskever, Timothy Lillicrap, Madeleine Leach, Koray Kavukcuoglu, Thore Graepel, and Demis Hassabis. Mastering the game of Go with deep neural networks and tree search. Nature, 529(7587):484-49) 2016.
Aleksandar Stanic, Yujin Tang, David Ha, and Jürgen Schmidhuber. Learning to Generalize With Object-centric Agents in the Open World Survival Game Crafter. IEEE Transactions on Games, 16(2):384-395, 2024.
Simon Sun. Monpachann-hollookkkiggz-zoo January 2024. URL https://github. com/
R.S. Sutton and A.G. Barto. Reinforcement Learning: An Introduction. The MIT Press, second edition, 2018.
TeamCherry. Hollow Knight, 2017. URL
Unity Technologies. Unity Documentation, 2005. URL https://docs...ee uni.yyc...ity.comm
Mark Towers, Jordan K. Terry, Ariel Kwiatkowski, John U. Balis, Gianluca de Cola, Tristan Deleu, Manuel Goulão, Andreas Kallinteris, Arjun KG, Markus Krimmel, Rodrigo Vicente, An- drea Pierré, Sander Schulhoff. Jun Jet Tai, Andrew Tan Jin Shen, and Omar G. Younis. Gymna- sium, March 2023. htss::::: zenodo. org/ waarrg//ecoooo rwcore/cooieeeee 81 27025
Ruben S. van Bergen and Pablo Lanillos. Object-hased Active Inference. In International Workshop on Active Inference, volume 1721 of Communications in Computer and Information Science, pp- 50 60.64. Springer, 2022.
Aäron van den Oord, Oriol Vinyals, and Koray Kavukcuoglu. Neural Discrete Representation Learn- ing. In Advances in Neural Information Processing Systems 30: Annual Conference Neural Information Processing Systems 2017, December 4-9, pp. 6306-6315, 2017.
Rishi Veerapaneni, John D. Co-Reyes, Michael Chang, Michael Janner, Chelsea Finn. Jiajun Wu, Joshua B. Tenenbaum, and Sergey Levine. Entity Abstraction in Visual Model-based Reinforce- ment Learning. In Annual Conference on Robot Learning, volume 100 of Proceedings of Machine Learning Research, pp. 1439-1456. PMLR, 2019.
Qianqian Wang, Yen-Yu Chang, Ruojin Cai, Zhengqi Li, Bharath Hariharan, Aleksander Holynski, and Noah Snavely. Tracking Everything Everywhere All at Once. In IEEE/CVF International Conference on Computer Vision, pp. 19738-19749, IEEE, 2023.
Nicholas Watters, Loic Matthey, Sebastian Borgeaud, Rishabh Kabra, and Alexander Ler- chner. googlc-deepmin/sspriooooll June 2024. URL atttps://gitie com/
ge hetggleeedeemmmmmn/pppeeeeeeoooe
Fandom Wiki. Category:Charms Hollow Knight Wiki, 2018. URL URLLttps::/hollowwniithh
Fandom. com. Sttmm//iiii/Caaegee Charms
Haofei Xu, Jing Zhang, Jianfei Cai, Hamid Rezatofighi, Fisher Yu, Dacheng Tao, and Andreas Geiger. Unifying Flow, Stereo and Depth Estimation. IEEE Transactions on Pattern Analysis and Machine Intelligence,
2023.
13

---

Published as a conference paper at ICLR 2026
Jinyu Yang, Mingqi Gao, Zhe Li, Shang Gao, Fangjing Wang, and Feng Zheng. Track Anything: Segment Anything Meets Videos. CoRR, abs/2304.11968, 2023.
Zhantao Yang. secrme//HollowKnnggh..tt 2023.
URL
http://tttuub..co/seeeeeeee
Holl Hollowknight_Bi RL
Denis Yarats, Ilya Kostrikov, and Rob Fergus. Image Augmentation All You Need: Regular- izing Deep Reinforcement Learning from Pixels. In 9th International Conference on Learning Representations, 2021.
Denis Yarats, Rob Fergus, Alessandro Lazaric, and Lerrel Pinto. Mastering Visual Continuous
Control: Improved Data-augmented Reinforcement Learning. In International Conference Learning Representations, 2022.
Weirui Ye, Shaohuai Liu, Thanard Kurutach, Pieter Abbeel, and Yang Gao. Mastering Atari Games with Limited Data. In Advances in Neural Information Processing Systems 34: Annual Confer- ence on Neural Information Processing Systems 2021, NeurIPS 2021, pp. 25476-25488, 2021.
Jaesik Yoon, Yi-Fu Wu, Heechul Bae, and Sungjin Ahn. An Investigation into Pre-training Object- centric Representations for Reinforcement Learning. International Conference on Machine Learning. volume 202 of Proceedings of Machine Learning Research, pp. 40147-40174, PMLR, 2023.
Tianhe Yu, Deirdre Quillen, Zhanpeng He, Ryan Julian, Karol Hausman, Chelsea Finn, and Sergey Levine. Meta-World: A Benchmark and Evaluation for Multi-task and Meta Reinforcement Learning. In Annual Conference on Robot Learning, volume 100 of Proceedings of Machine Learning Research, pp. 1094-1100. PMLR, 2019.
Matthew Zeiler, Dilip Krishnan, Graham W. Taylor, and Robert Fergus. Deconvolutional net- works. In IEEE Conference on Computer Vision and Pattern Recognition, pp. 1528-25335 IEEE Computer Society, 2010.
Renrui Zhang. Zhengkai Jiang. Ziyu Guo, Shilin Yan, Junting Pan, Hao Dong. Yu Qiao, Peng Gao, and Hongsheng Li. Personalize Segment Anything Model with One Shot. In International Con- ference on Learning Representations, 2024.
Weipu Zhang, Gang Wang, Jian Sun, Yetian Yuan, and Gao Huang. STORM: Efficient Stochastic Transformer based World Models for Reinforcement Learning. In Advances in Neural Informa- fion Processing Systems, 2023.
14

---

Published as a conference paper at ICLR 2026
A DETAILED MODEL STRUCTURE
Spetil
Pradiced test sxps" sledhate aaora [L. 16 or
Predicedd lighs reword and lemeration IL 253] and II
Terpocal an - FFN
No Micko
Coeas ann. acruss Hirourps
MI Prodicsor
prodictor
Speiial sits FFN Self across obeess
Hikken stace [I. K. 256)
Aqion
Trainable positood enbollii [. 156)
Sprisitzmprrl
Pollcy
[L 25g]
Value
MIP
Object rachoeic lexzte [I. 16-16]
Mieai stacunii lakens
Stecholi Leceb.
Artion (c.,Action din)
[L. 3Dx32]
KL, 16-16 S2x31
enogerial
ettegrredd
Crogorica VAE
om predkter
Objar loters K, 2wes
Resiad eharaii (L,3 $4,54)
stxchatic litett
Calquriccl
Lioor
Straght-throsph gradient
sagge
Olocrvatica
Escoda:
Deooder
Stf
[, HL W]
layets
Reconannations
Inceable
licilen
Ofeect madale
Viaaal malak
258]
Figure 5: The model structure of our proposed OC-STORM. The tuples in square brackets represent the shapes of the corresponding tensors, where L denotes the batch length or sequence length, K is the number of objects, and H and W are the image height and width, respectively. The object module constitutes the proposed object-centric component, while the visual module processes re- sized raw observations. K* equals K when only object module is used, equals when only visual module is used, and equals + when both modules are used. The trainable token and positional embeddings are broadcast to match the shapes of the corresponding tensors. The reward logit is 255-dimensional and used for the symlog two-hot loss (Hafner et al. 2023).
B LOss FUNCTIONS
B.1 WORLD MODEL LEARNING
The world model is trained in a self-supervised manner, optimizing it end-to-end. Our setup closely follows D3eamerV3 (Hafner al. 2023), with details presented for completeness. We use mean squared error (MSE) loss for reconstructing the original inputs, symlog two-hot loss Laym (Hafner jet al., 2023) for reward prediction, and binary cross-entropy (BCE) loss for termination signal pre- diction. These losses, collectively referred to prediction loss, are defined as:
pred (Ø)
||se sal1
Laym (fa Te log f. (1 - Te) log(1 ft)
(5)
Recunstruction Juss
Reward loss
Termination loss
The dynamics loss cdym (Ø) guides the sequence model in predicting the next distribution. The representation loss (() allows the encoder's output to be weakly influenced by the sequence model's prediction, ensuring that the dynamics are not overly difficult to learn. These losses are identical Kullback Kullback-Leible Leibler (KL) divergence losses except for their gradient propagation settings. We use sg(-) to denote the stop gradient operation. The dynamics and representation losses are defined as:
Ldyu(d) = max (1, KL
|| 9 dvm (b)41)
(6a)
(o) max(1, (1, KL (2441114441
sg(g Ayn (h.)))
(6b)
The max operation represents free bits for KL divergence, encouraging the model to focus on opti- mizing prediction losses for better feature extraction the KL divergence is too small.
The total loss function for training the world model is calculated as follows, where Ep denotes the expectation over samples from the replay buffer:
((w) = ED Cproa (0) + Layn (0) + 0.5
(7)
15

---

Published as a conference paper at ICLR 2026
The coefficient of 0.5 for Crep is used to prevent posterior collapse (Lucas et al. 2019 situation where the model produces the same distribution for different inputs, causing the dynamics loss to trivially converge to O. The imbalanced KL divergence loss helps to mitigate this issue.
B.2 POLICY LEARNING
The policy learning approach closely follows that of DreamerV3 (Hafner et al. 2023). with modi- fications specific our method. The key differences lie in the input for the policy and the action dimension for the game Hollow Knight. We use the concatenation of object latents, object hidden states, visual latents, and visual hidden states input features. For Hollow Knight, a multi-discrete action space is employed.
The agent learns entirely on the imagined trajectories generated by the world model. To begin the imagination process, we first sample short contextual trajectory from the replay buffer. During imagination, future environmental inputs 1: are unknown, and sampling from posterior dis- tribution (0(=x(5)) is unavailable. Thus, we sample the latent variable from the prior distribution 96 90 "E (he) and optimize the policy over êt+1. However, during testing, the agent interacts di- rectly with the environment, allowing access to the posterior distribution of the last observation. This introduces a difference in notation. For simplicity, we do not distinguish between and & in the following descriptions.
The agent uses both the latent variable & and hidden states he as inputs, as defined below:
Critic:
Vo he) ) se Exo.o Σ rree
(8)
k=D
Actor:
a ~ To le; h4)-
Here, 7' is the number of timesteps in the episode. We use two separate MLPs for the critic and actor networks. The symbol & indicates that the trajectories are generated within the imagination process of the world model.
For value loss, we employ the A-retum G (Sutton & Barto value estimation. It is recursively defined as follows, where is the reward predicted by the world model, and ft represents the predicted termination signal:
2018 Hafner et 2023) improve
G = ft + 7(1 +x) (1
+ {
(9a)
G = (=1,
(9b)
To regularize the value function, we maintain an exponential moving average (EMA) of the critic's parameters, defined in Equation equation [0] This regularization technique stabilizes training and helps prevent overfitting, where v/, represents the current critic parameters, a is the decay rate, and EMA denotes the updated critic parameters: *++
EMA
ov. EMA + (1- o)v
(10)
++1
For policy gradient loss, apply return-based normalization for the advantage value. The normal- ization ratio S is defined in Equation equation Taas as the range between the 95th and 5th percentiles of the X-return G across the batch (Hafner al. 2023):
S percentile 95) percentileeG?..o.
(11)
The complete loss functions for the actor-critic algorithm are given by Equation equation 12]
- Vv(ss) In To(arlz h.) gH HH(Hn (To(oole.,))
[(0) = Ema.o
-sg
m) (
(12a)
(Vo(zr, sgg.:hh))ssCi.. sg((:))))))))))>>>>>))) Luym (Vo(ze, he), sg (V (zt, (zt, h.
((()) = Emo.o
Luym
(12b)
Here, H.)) denotes the entropy of the policy distribution, and n7 = - is the coefficient for entropy loss.
16

---

Published as a conference paper at ICLR 2026
C COMPARISON WITH OTHER MBRL ALGORITHMS
Table 2: Game scores, overall human- normanno scores, and computing time on the Atari 100k benchmark. Bold indicates scores within 5% of best. The last column '#obj' denotes manually annotated objects per game. The Dreamer V3 scores here are the original data from their paper (Hafner al. ).23). SPR (Schwarzer al. 2021) is typical sample-efficient model-free method. IRIS (Micheli et al. 2023) and TWM (Robine et al. 2023) are transformer-bas based MBRL approaches,
and A-IRIS (Micheli al., 2024) is an upgraded version. DIAMOND (Alonso al. 2024 is
diffusion-baat baseed MBRL approach.
Game
Random
Human
SPR
IRIS
TWM
DreamerV3
A-IRIS
DIAMOND
OC-STORM
Alicn
228
7128
842
420
675
1278
599
744
1101
Amidar
1720
120
226
163
Assault
222
742
566
1524
683
741
1435
1526
1270
Asterix
210
8503
1117
1020
2001
3698
1754
BankHeist
753
345
53
467
422
1206
20
1075
BattleZone
2360
37188
14834
13074
5068
20800
10365
4702
4590
Boxing
Breakout
2
30
20
84
20
11
226
132
52
ChopperCommand
811
7388
1565
1697
2440
1101
1370
2090
CrazyClimber
10780
35829
36700
59324
71820
80060
70920
99168
84111
DemonAltack
152
1971
518
2034
350
454
884
288
411
Freeway
Frostbite
65
4335
1171
259
1476
3914
287
274
260
Gopher
258
2412
661
2236
1675
2252
9349
5898
4457
Hero
1027
30826
5859
7037
7254
13324
6235
5622
6441
Jamesbond
29
303
366
463
362
490
345
427
347
Kangaroo
3617
Krull
1598
2666
3682
6616
6349
8604
6392
8610
9715
KunstuMaster
258
22736
14783
21760
24555
25560
251559
18714
24988
MsPacman
307
6952
1318
999
1588
1400
1175
1958
2401
Pong
-21
15
-5
15
19
-5
15
20
21
PrivateEye
86
100
Qbert
164
13455
866
746
3331
3700
3438
4499
4546
RoadRunner
7845
12213
9615
9109
18440
10622
20673
20482
Seaquest
42055
558
661
774
964
895
551
712
UpNDown
533
11693
10859
3546
15982
49456
8091
3856
6623
HNS mean
0%
100%
61.6%
104.6%
95.6%
128.2%
138.3%
145.9%
134.8%
HNS median
0%
100%
39.6%
28.9%
50.5%
48.7%
59.4%
37.3%
43.8%
Average compute
Py Torch
PyTorch
PyTorch
Jax 0.1
PyTorch 1.0
PyTorch 2.9
PyTorch 0.2
4090 days
0.1
17

---

Published as a conference paper at ICLR 2026
D REVIEW OF OBJECT REPRESENTATION METHODS
Table brief review of state-of-the-ar methods in different fields of computer vision related to object-centric reinforcement learning.
Method
Description
Cutic
Retrieval-based semi-supervised vidco object segmentation algorithm. Provides compact vector representation of objects using an object transformer (Cheng
at. 2023].
XMem
Retrieval-based -emi-supervised vidco object segmentation algorithm with multi-level memory system. Lacks compact object representation (Cheng &
Schwing 2022].
SAM
Open-sct image segmentation algorithm. Generates masks via user prompts but requires prompt for each frame in video (Kirillov et al. 2003).
TrackAnything
SAM XMem for finc- grained vidco segmentation, Requires twice the compu-
tational resources compared to XMem (Yang ct al. [2023).
PerSAM
Onc-shot enhancement of SAM. Difficult expand to few-shot cases (Zangg
al. 2024).
SAM2
Statc-of-thc-art multi-usage vidco object segmentation algorithm. Capablc of providing compact vector representations. (Ravi et al. (2024).
YOLO
Closed-set object detection algorithm. Requires extensive annotations for train-
ing or fine-tuning (Redmon ct al. 2016] ct al. [023).
GroundingDINO
Open-set object detection algorithm. Generates object bounding boxes using natural language prompts but struggles with rare or abstract cases, such as video game players (Liu ct al. (2023).
Slot attention
Unsupervised image object discovery and segmentation algorithm. Underper- forms compared to supervised methods (Locatello al. 2020)
SAVi
Unsupervised video object discovery and segmentation algorithm. Evaluation is
semi-supervised, though training is unsupervised. Underperforms compared semi-supervised methods (Kipf ct al. 2022
( Elsayed et al., [2022].
Omnimotion
Video point-tracking algorithm. Provides pixel-level tracking but lacks object-
level information and does not support few-shot scenarios (Wang al. 2023
Unimatch
Dense optical flow and depth estimation algorithm. Useful for moving identity
detection but cannot extract object-level information and struggles with out-of- domain generalization (Xu al. 2003)
E HOLLOW KNIGHT
E. NUMERICAL RESULTS
Table 4: Episode returns and win rates (WR) of STORM and the proposed OC-STORM Hollow Knight. Each seed's performance is measured by the mean episode return across 20 runs, and the average of these 3 mean returns is reported. The "#obj" column shows the number of annotated objects for a boss. Scores that are the highest or within 5% of the highest score are highlighted in bold. To evaluate the upper limit of our agent, conduct 400k run Pure Vessel, which demonstrates that our agent can defeat one of the most difficult bosses in the game with sufficient training.
Boss name
Random
Optimal
STORM
OC-STORM
STORM WR
OC-STORM WR
#nbj
God Tamer
56
35.0
41.7
70.0%
55.0%
Hornet Protector
37
28.1
32.4
66.7%
100.0%
Mage Lord
38
19.6
28.0
5.0%
48.0%
3
Mantis Lords
42
33.2
35.2
71.7%
83.3
3
Mawlek
41
36.9
98.3%
98.3%
Pure Vessel
55
8.9
15.7
0.0%
0.0%
Pure Vessel (400k)
2
55
25,3
35.0
6.7%
13.3%
2
18

---

Published as a conference paper at ICLR 2026
E.2
RELATED WORK
Despite its popularity among players, Hollow Knight has seen limited use as a benchmark for re- search in reinforcement learning. We introduce repositories, published research, and other relevant resources that leverage or explore Hollow Knight as benchmark. Cui (2021) employs DQN (Mnih at 2013) and its variants but requires modding the game background to black enhance character perception. Yangl (2023) uses the Rainbow algorithm (Hessel et al. 2018) with additional techniques like DrQ (Yarats et al. 2022 202T), achieving high win rates against several of the game's bosses.
Yang's repository has been widely forked and adopted. Building on his work, [Lee (2023) studies the effect of reward shaping. while [Sun| (2024) focuses on improving training efficiency by tuning the game interaction configuration and switching to the PPO algorithm (Schulman et al. 2017). [Jain (2024) leverages internal game states to extract hitboxes as input for the algorithm, representing them as segmentation masks that passed to DQN or PPO.
E.3 ENVIRONMENT CONFIGURATION
Hollow Knight is a modern video game developed with Unity (Technologies, 2005). To our knowl- edge, efficient simulators for this game, such as those available for Atari (Brockman et al.) [016 Towwess et al. 2023), are not publicly available. Therefore, we developed a custom wrapper that captures screenshots of the game 9 FPS and sends keyboard signals execute actions. The FPS rate is a choice based on the author's experience with the game and considerations for computational efficiency. To obtain reward signals, we developed a modding plugin (Bham & Wyza (017) that logs when the player-controlled character (the Knight) either hits an enemy or is hit. Our wrapper then parses this log file to generate reward and termination signals.
The game execution and agent training are conducted on Windows machine. To monitor training progress and statisties without interrupting the game, we needed method send keyboard inputs to a background or unfocused window. However, Windows lacks an API for this purpose. As a result, the game window must remain in the foreground, fully occupying the training device and hindering monitoring. To address this, utilized Hyper-V V (Cooley) 2022) Windows virtual machine to run the game in the background, with Ray (Moritz et al.] (018)) facilitating communication between the host and virtual machine. Training and processing occur on the host machine, while the virtual machine handles interactions with the environment. This setup can be extended to distributed nodes, with some handling game rendering and others managing training tasks.
For in-game configuration, the charms (Wiki 2018) are set to Unbreakable Strength, Quick Slash, Soul Catcher, and Shaman Stone across all experiments. This configuration is chosen to explore the agent's fighting potential rather than its gliteh-finding abilities.
E.4 ACTION SPACE
All previous works utilize human-specified action space rather than the original keyboard inputs. For example, in the Yang (2023) implementation, short and long jumps are treated as two distinct actions, which are originally controlled by the duration of the jump button press. His environment wrapper handles this difference with a fixed command. While this design reduces the exploration and computational costs for reinforcement learning agents, it cannot capture the full range of possible actions in the game. Advanced operations, as demonstrated in this video (Cranky Templar 2018), require full control of the keyboard. Therefore, we design the action space as a multi- multi-binary--eeeeee binarr-ddaceeeeettt one that directly binds to the press and release of the physical keyboard, which will be explained in Section E.4
We design the action space as a culti--binary-disceed space directly mapped to the press and release states of eight specific keys on the keyboard. These keys include W, A, S, and D for movement directions, and J, K, L, and I for attack, jump, dash, and spell actions, respectively. Each key's state is represented as a binary variable, where corresponds to a key release and 1 corresponds to key press. The action can therefore be described the binary state of one key. The key's state is maintained between frames, and a toggle signal is sent only when there is change in the key state from 0 (press) 1 0 (release).
{0. 138, with each element representing
19

---

Published as a conference paper at ICLR 2026
The probability of an action is determined by the independent probabilities of each key's state:
T0) no ao (2 h))
(13)
where at denotes the state of the k-th key.
The entropy of the action space, H (TT(al2,, hi). is the sum of the entropies of the individual key states:
( h))
Σ H(oo(a
(14)
This design provides fine-grained control over the agent's actions, allowing for the execution of complex manocuvres while maintaining tractable exploration space for reinforcement learning,
E.5 REWARD SHAPING HOLLOW KNIGHT
Most existing methods (Yang 2023 Sun 2024 Jain 2024 Cui 2021
Lee 2023) for Hollow Knight
use reward structure of +I for hitting an enemy and for taking damage. Some approaches mod- ify the weighting ratios, while others introduce auxiliary rewards for performing specific actions. However, we found that these settings are suboptimal for training reinforcement learning agents.
Our method assigns a +1 reward signal for hitting an enemy and virtual termination signal upon being hit. The game continues until the episode naturally ends. The termination signal is stored in the replay buffer for training the world model, treating health loss as a life-loss event. Leveraging lose-loss information is a common technique that aids in value estimation YYe et al., 202T] Micheli et al. 2023] Zhang et al. [023)] Alonso et al. (024). Additionally, the Knight can damage enemies in multiple ways, and these damages are normalized against the base attack damage to compute the positive reward.
HornePrrotector
MantisLonk
40
50-
20
10
50k
10ck
50k
100G
OC SCC-SORM
Logacy rewards
Figure 6: Training episode returns for Hollow Knight's Hornet Protector and Mantis Lords under different reward settings, "Legacy rewards" refer to the reward scheme used in prior works. For comparison, we aligned the returns from "legacy rewards" with our baseline settings by accounting for lost health.
Here, we present the key differences between the two reward settings. As illustrated in Figure 6 our reward configuration is more robust than those used in previous studies, resulting in signifi- cantly improved performance, especially in more challenging environments like Mantis Lords. This improvement can be analyzed from two perspectives:
1. Terminating the episode upon being hit better aligns with human cognition and the agent's s expected behaviour. The aim is for the agent deal as much damage as possible without taking any. While this may seem aggressive, raising concerns that the agent might sacrifice itself to deal more damage, neither our qualitative nor quantitative results show this ten- dency. Survival naturally offers more opportunities deal with future damage, which the agent learns to prioritize. Although applying a negative penalty for being hit could prevent the agent from sustaining multiple consecutive hits in highly unfavourable situations, such scenarios should not occur under an optimal or near-optimal policy.
2. While maintaining the same optimal policy, truncating future rewards upon being hit sig- nificantly reduces the variance in value estimation. Hollow Knight is highly stochastic
20

---

Published as a conference paper at ICLR 2026
environment where bosses behave aggressively yet unpredictably. Estimating value directly over an episode (lasting approximately 300 700 timesteps) is inherently challenging in such settings.
E.6 COMPARISON WITH MODEL-FREE BASELINE
As introduced in Section E.2 Yang's repository YYng 2023 is widely recognized implementa-
tion within the community. this section, we compare the performance against the boss Hornet Protector.
Yang's reward structure assigns +0.8 for hitting an enemy and -0.8 for taking damage, with addi- tional auxiliary rewards on the order of x 10-4 for various actions. A small feedback reward is also given at the end each episode. The choice 0.8 weight factor for rewards reflects the use of +1/-1 reward clipping, with margin reserved for the auxiliary rewards. We provide broad comparison with this approach below.
HornetProtector episode return
HameeProecctor win rate
20
0.5
1006
rewards
Yang
Figure 7: Training episode returns and win rates on Hollow Knight's Hornet Protector with our proposed method and Yang's (Yang) 2023) method. "Legacy rewards" are described in Section
E.3 We applied some preprocessing to align the two returns for easier comparison, so the "legacy rewards" curve may appear different from the one shown in the previous section. The win rate is more straightforward and can be used for comparison without changing.
As shown in Figure 7 our implementation is more efficient than Yang's s. As we noted in Section 4.2 there are significant differences between our methods, making this not strictly a fair comparison from an algorithmic standpoint. This comparison intended solely to demonstrate the efficiency of our implementation.
Additionally, Yang claims that his agent can achieve 10 wins out of 10 battles, which is accurate de- spite his win rate in our plot appearing to be lower than 100%. Two reasons may lead to this. First, his original sample steps are greater than ours, which may account for differences in performance. Second, our in- na-g charm configuration (Wiki) 2018) differs from the used in his implemen- tation. When testing Yang's 's implementation, we retained our current charm settings, which likely impacted the win rate results.
21

---

Published as a conference paper at ICLR 2026
F ADDITIONAL AND ANALYSIS
F.1 ATTENTION-BASEE POLICY OR MLP-BASED POLICY
When handling multiple objects, an iitention-baseld policy network like the self-attention predictor described in Appendix iA natural choice. A previous work OC-SA (Stanic et al. 2024) has also explored such structure. However, we still design our actor and critic networks MLPs which take the concatenation of object latent variables and hidden states as input.
We found that the brest policy tends to overfit pre-learned behaviours and makes it hard to learn new knowledge. This will not be a major issue in stationary games like Boxing but will face trouble in non-stationary games like Pong. For example, the attention-basee based policy can quickly learn how to catch the ball but cannot efficiently learn how to score against the opponent. On the one hand, can confirm this by visually checking the rendered episodes. On the other hand, numerically speaking, we can observe the episode length of playing Pong. If the episode length increases while the episode returns remain at the same level, then we can tell that the agent learns how to catch the ball, but is stuck in that local optimum.
Baxing
Alari Forg: spisode length
3000
50-
2020
20
10k 1ock
100k
10CK
MI.P-bosrd policy
Attention-hoeed policy
Figure 8: Training episode returns for Atari Boxing, Pong and episode lengths for Pong of attention- based policy and MLP-based policy. The attention based policy can learn as quickly as the MLP- based policy for catching the ball but struggles to transition to the scoring phase in Atari Pong.
As the results plotted in Figure 8 can tell that the attention-based policy suffers from that issuc. The episode length of both policies rises at a similar speed before 50k steps, but it declines slower for the attention-based policy after that. The experiments are conducted using only the object module, so the MLP-based policy curves are identical to the "vector" ones in Figure 4al As the visual latent itself contains all the information, the agent can choose only use that part of the information and thus may affect our judgement on the effectiveness of the attention-based policy.
Though the attention-based policy has the potential to handle a dynamic number of objects, our experiments are conducted on fixed number. As it doesn't demonstrate superior performance than the MLsed policy our case, we always use MLP-based in other tasks for consistency in evaluation.
F.2
IMPACT OF THE NUMBER OF ANNOTATIONS
Since Cutie is a retrieval-basee algorithm that stores past frames and masks in a buffer for reference, naturally supports the use of multiple annotation masks beyond the first frame by substituting model-generatee generated masks in the buffer. Incorporating more label masks can capture a wider range of object states, leading to more consistent segmentation results. However, reducing the number of labels can further lower annotation costs and computational complexity. In this section, we explore the impact of the number of labels on agent performance. To eliminate the influence of visual input, conduct these experiments using only the object module.
As shown in Figure increasing the number of annotation segmentation masks enhances the robust- ness of the agent's 's performance, even in visually static environments like Atari Boxing and Pong. In these environments, a single frame can include all necessary objects for decision-making However, Cutie may lose track of objects if their states deviate significantly from those in the labelled masks, such as the punching state versus the standing state in Boxing, or the paddle in Pong when it is partially off-screen and appears shorter than when centred.
22

---

Published as a conference paper at ICLR 2026
Boxing
Pong
16-
-20
9M
100k
53k
1006
6-shot
I-shot
Figure 9: Training episode returns for Atari Boxing and Pong with different numbers of annotation segmentation masks. Increasing the number of annotation masks enhances the robustness of the agent's performance.
Moreover, in complex environments like Hollow Knight and Minecraft, a single frame may capture all objects, which often necessitates additional segmentation masks. For consistency in evaluation, we use six annotation segmentation masks for Atari and twelve for Hollow Knight.
F.3 THE COMPUTATIONAL OVERHEAD OF OC-STORM
The computational overhead of OC-STORM on Atari games with an NVIDIA GeForee RTX 3090 is shown in Table 5 The input resolution for Cutie is 420 x 320 (double the original 210 160).
Thus the computational cost of introducing Cutie acceptable in many cases.
Table 5: The computational overhead of OC-STORM on Atari games. The three numbers in block here mean sample evaluation speed (iterations/secondd), training speed (iterations/second) and hours spent for train with a 100k sample budget, respectively.
Algorithm
objects
object
objects
3 objects
STORM
114 it/s
8.1 itls
3.67
(visual module only)
OC-STORM
32 it/s
8.8 itss
4.02 h
32 it/s
8.5 itls
4.14
31 itss
7.8 it's
4.46
(obj module only)
OC-STORM
28 it/s
5.9 it/s
5.70
27 it/s
5.5 it/s
6.08 h
27 it/s
5.3 it's
6,27 h
(both modules)
We report the latency of each part in our pipeline in T As shown, the introduced overhead is modest, and the system maintains real-time performance throughout all evaluated settings.
Table 6:
Average system-level runtime (ms) of the proposed model on a single NVIDIA RTX 4090
GPU, with segmentation inputs at resolution of 420x320 pixels. The 10 objects case uses manually constructed synthetic configuration for testing.
Module
objects 3 objects
4 objjcts
10 objects
Cutie Load
1500
Sample step
15
15
16
21
Segmentation Inference
11
11
12
15
Policy Execution
4
4
4
6
Train step
180
212
332
World Model Gradient Step
41
World Model Imagination
113
142
257
First Step Imagination
4.5
5.3
11.0
Last Step Imagination (16 steps)
5.6
13.4
Policy Grudient Step
28
23

---

Published as a conference paper at ICLR 2026
G DETAILS FOR THE USE OF SAM2 AND CUTIE
G.1 NUMBER OF ANNOTATIONS, INPUT RESOLUTIONS, AND MODEL SIZE
To prompt Cutie, we use annotation masks per Atari game and 12 per Hollow Knight boss. One potential critique is that few-shot annotation requires prior knowledge of the environment, which may seem unsuitable for general agent learning. However, we view this process as akin inform- ing the agent of certain task rules. While rewards can reflect task rules, they often too sparse to facilitate an understanding of complex environments. Just as humans may initially struggle to understand how to play game without being told the rules, there is reason not inform agents of key objects. Therefore, we believe this pipeline holds practical value in many cases.
For Atari, we upscale the observation from 210 X 160 420 320. This upscaling aids in the identification of small objects in Atari games, such as the ball in Pong and Breakout. For each game, we hand annotate 6 masks. For Hollow Knight, resize the observation's shorter side to 480p while maintaining the aspect ratio before inputting it into the Cutie. For each game, we hand annotate 12 masks.
For Cutie, we use cut emaamal 1-mega. pth;; for SAM2, we use sam2. hiera small.pa along with the corresponding configurations, respectively.
G.2 MODIFICATIONS FOR INTEGRATION WITH STORM
We make no modifications to the official implementation, except for caching and copying internal variables. The only special process involves setting the object feature vector to when the model loses track of the object. This can prevent the world model from fitting random dynamics in such cases.
For SAM2, this is rather straightforward since the model naturally contains model to predict whether the target occluded, and we can directly use this score as the determinant.
For Cutie, the model uses attention guidance mask within its object transformer, which restricts which visual features the object feature can attend to. This mask is trained as part of an auxiliary segmentation task. When the attention guidance mask is set to all Is (0 allows attention and rejects it), indicating that Cutie cannot find strong evidence of the object's presence in the scene, the transformer theoretically should reject all attention from the visual features.
However, in this situation, Cutie inverts the mask, allowing the object feature to attend to all visual features in an attempt search for the object in the scene. As result, the attention becomes scattered across the observation space, leading to unpredictable output for the object feature. Thus, we set the object feature vector when the attention guidance mask is entirely
24

---

Published as a conference paper at ICLR 2026
H HYPERPARAMETERS
Table Hyperparameters for both Atari and Hollow Knight. The life loss information configuration aligns with the setup used in EfficientZero YYe al. 2021). Regarding data sampling. each time sample B1 trajectories of length for world model training, and sample B2 trajectories of length C for starting the imagination process. The train ratio is defined the number of gradient steps over the number of environment steps.
Hyperparameter
Symbol
Value
Transformer layers
K
2
Transformer feature dimension
D
256
Transformer heads
4
Dropout probability
p
0.1
World model training batch size
B
32
World model training batch length
T
32
Imagination batch size
B2
512
Imagination context length
C
4
Imagination horizon
L
16
Train ratio
Environment context length
16
Gamma
0.985
Lambda
0.95
Entropy coefficiency
1
1 10-3
Critic EMA decay
0
0.98
Optimizer Activation functions
Adam Kingma & Ba
2015)
SiLU
Flfwing et
2018)
World model learning rate
1.0 x TO 1000
World model gradient clipping
Actor-critic learning rate
3.0 10-5
Actor-critic gradient clipping
100
Gray scale input
False
Frame stacking
False
Atari frame skipping
4 (max over last 2 frames)
Hollow Knight FPS
Use of life information
True
25

---

Published as a conference paper at ICLR 2026
I ILLUSTRATION OF LIMITATIONS
Miss lord.
Successfully identify lords.
Figure 10: Sample frame and segmentation masks generated by Cutie from the Hollow Knight Mantis Lords. Cutie may lose track of one of the lords (represented with green masks). This tracking issue more likely to occur not only in this scenario but also in other environments where duplicated instances are present, compared to scenes with a single instance.
100
(a) Sample frame.
(b) Annotation mask.
Figure 11: Sample frame and annotation segmentation masks from Atari Gopher. We only specify two objects for Gopher. The tunnel in the ground, which is the navigable space for the gopher, is challenging encode as object given model structure.
26

---

Published as a conference paper at ICLR 2026
J SAMPLE ANNOTATIONS FOR ATARI, HOLLOW KNIGHT, AND METAWORLD
and 14 present sample frames and annotations used by our method in Atari, Hollow Knight, and Metaworld, respectively. For each task in Atari and Metaworld, we annotate 6 frames, and for each boss in Hollow Knight, we annotate 12 frames.
Figure 13
Observation
Annotation mask
HomeiPPotecch
Magel Lond
MuntisLords
Mawi
Vessel
Figure 12: Sample frames and annotation masks for Hollow Knight bosses.
Boxing
Breakunt
Gupber
KumeTa Master
Pony
Obert
RoadRumner
Observation
Armocation mask
f.1
Figure 13: Sample frames and annotation masks for Atari games.
27

---

Published as a conference paper at ICLR 2026
basketball pick-out-ffhhlle
Observation
Asmolation mask
Figure 14: Sample frames and annotation masks for Meta-world tasks.
K ABLATION ON MISSING OBJECTS
We conducted ablation study on removing annotated objects, and all experiments only object features with no visual input for the world model. We found that:
Missing the object that directly influenced or controlled by the input actions causes significant performance drop. Increasing the number of tracked objects generally improves performance.
2. Some objects (e.g., the brick wall in Breakout) may negatively affect performance. We believe this arises from current video segmentation algorithms' limited sensitivity to subtle object-level changes (the entire brick wall treated as single object). This limitation discussed in Appendix
Boxing
Breakout
30-
30
-30
24
100%
Figure 15: Ablation study on missing objects. Sample game observations can be found in Figure 13
28

---

Published as a conference paper at ICLR 2026
L UNCONDITIONAL SAM ON HOLLOW KNIGHT
At the beginning of our study, we also explored the possibility of performing segmentation using SAM (Kirillov al. 2023 Ravi al. (024) without any labeled frames. The results, as illustrated in Figure 16 show that SAM can roughly distinguish different objects in the Hollow Knight envi- ronment; however, the alignment of semantic meaning and the delineation of object scales were not satisfactory in the absence of prompts. As result, the quality of the object feature wouldn't be sufficient support the control objectives of our framework. These observations motivated to incorporate targeted supervision into our method design.
61,00
6890
Figure 16: Sample unconditional segmentation results with SAM on Hollow Knight.
M THE USE OF LARGE LANGUAGE MODELS (LLMs)
We use LLMs for polishing and grammar checking. No content is directly generated with LLMs.
29


Reviews:
[{"review_id": "yuj9fZJO0e", "reviewer": "Reviewer_FttL", "summary": "The paper proposes OC-STORM, an object-centric model-based RL framework that fuses few-shot, pretrained video-segmentation features with pixel inputs to train a spatial-temporal world model, yielding improved sample efficiency on Atari-100k and Hollow Knight.", "questions": "- Labeling cost & K sensitivity. How many minutes of annotation per game are required in practice, and how does performance vary with K (under-/over-specifying the number of tracked objects)? Could you add a curve for returns vs. K and vs. number of annotated frames? \n\n- External baselines & reporting. Can you include a matched-config comparison against recent token/diffusion world models and standardized Hollow Knight setups (or release your wrapper to make one), plus runtime/latency tables for the segmentation pipeline?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "- The method integrates few-shot object features from SAM2/Cutie into a spatial–temporal world model (Transformer/RNN backbones), with clean modality separation (object tokens + visual token) and categorical VAE discretization; the training and architecture are well specified. \n\n- On Atari-100k, object-centric variants outperform baselines (e.g., Cutie-OC-STORM reaches HNS mean 134.8% vs. STORM 114.2%, median 43.8% vs. 42.5%); Hollow Knight learning curves show faster convergence on harder bosses. \n\n- The paper diagnoses why vector features beat mask features, provides module ablations, feature-only reconstructions, and a failure-robustness study by zeroing object features; Meta-World results indicate portability beyond Atari.", "weaknesses": "- Comparative scope. Core comparisons are mainly within-framework ablations (STORM/DreamerV3 variants); external SOTA world-model baselines (e.g., diffusion/tokenization variants) and broader agent baselines are deferred or absent in the main text, and Hollow Knight lacks standardized settings—making cross-paper claims harder to calibrate. \n\n- Annotation/K configuration burden. The user-set K (objects) and handful of annotated frames (≈6–12) are reasonable but the human-time budget and sensitivity to K are not quantified in the main text; detection incompleteness is acknowledged but only coarsely analyzed. \n\n- System metrics incomplete. The paper references computational overhead analyses but does not report build/index/latency figures for segmentation + KG-like memory in the main body—useful for scaling to long episodes or high-res inputs.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "2FQfqCMjYh", "reviewer": "Reviewer_waoT", "summary": "Presents OC-STORM, an object centric model based RL methods:\n- The user gives the model a few annotated object masks which are used to get SAM2 / cutie object representations.\n- The object vectors and a downsampled image vector are embedded to a categorical latent (via categorical VAE)\n- A world model is trained via reconstruction\n- A policy is trained inside of the world model for control. \n\nPerformance results:\n- The object-centric variants of DreamerV3 and STORM tend to perform better than the non-object centric part. \n- Further, for games where SAM / Cutie is better at detecting objects in, performance is consistently better, which demonstrates the merits of the method when there is a good object extraction\n- Further evaluated in a hollow knight 100k setting and find that it out-performs the on-OC variant\n- Can also work in continuous tasks (Metaworld)\n- Analysis shows that (i) object information is well-learned as scene can be reconstructed from object vectors (players in boxing), and (ii) there is some robustness to randomly zero-ing out detected objects", "questions": "- How do you decide on how many key objects to label in the human annotated frames? Is there any guidance here for what to / not to label?\n- Can the segmentation model work without any labelled frames (I was under the impression that SAM can)? If so, how accurate are they at segmenting objects?\n- How sensitive is the policy to different human labels? If different users labelled the same game, or labelled a different number of objects, how would this change the segmentation and would the policy still be robust to this?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 2, "strengths": "The idea is novel in using few-shot labelled object masks to get (pre-trained) object representations, which is then used to train a world model for policy control.\n\nThe paper presents clearly and is of high quality. It presents positive results against reasonable baselines across a number of environments (atari100k, hollow knight, metaworld). It reads as a comprehensive work that can be informative for future works in object centric RL to build on.\n\nThe ablations are comprehensive. The authors do good science to isolate the effect of object encoding (e.g. object detectable vs. badly detected), if the object representation can be reconstructed, and robustness to segmentation model failure.\n\nFinally, I appreciate that authors are also frank in discussing some short-comings of current object extraction methods (such as inability to handle duplicate objects and geometric structures such as walls and floors) for future research to build on.", "weaknesses": "The major weakness of this method is the effort vs. gains trade-off for using this object-centric representation. To use OC-STORM, the user must first generate 6-12 frames of object mask labels for each environment they may wish to run. The gain from doing this, based on the paper, seems to be _mainly_ about better _sample efficiency_. One could argue that instead of going through the effort of labelling, the user can also (i) run the alternative methods longer to get similar performance, or (ii) use a more compute-intensive method to get the same sample efficiency (e.g. Delta-IRIS or DIAMOND in Appendix C seem to show this -- correct me if I mis-understood). \n\nOn the other hand, some of the original work on object-oriented RL were motivated by the hope that modelling object & interactions can allow generalization in fundamentally different ways [1]. More recent works have similarly shown zero-shot generalization [2,3], new ways of doing model learning & efficient exploration [3], and the ability for compositional generalization [4,5]. \n\nIf sample efficiency is indeed the main goal, and being mindful that future work can push this method to be even more performant than the more compute intensive methods, it still feels worthwhile to discuss the particular use case for the current approach as presented currently (for instance, in a low sample, low compute setting?), and how it may differ (or is similar) to other object-oriented approaches.  \n\n-----\n\n[1] Diuk, Carlos, Andre Cohen, and Michael L. Littman. \"An object-oriented representation for efficient reinforcement learning.\" Proceedings of the 25th international conference on Machine learning. 2008.\n\n[2] Sancaktar, Cansu, Sebastian Blaes, and Georg Martius. \"Curious exploration via structured world models yields zero-shot object manipulation.\" Advances in Neural Information Processing Systems 35 (2022): 24170-24183.\n\n[3] GX-Chen, Anthony, Kenneth Marino, and Rob Fergus. \"Efficient Exploration and Discriminative World Model Learning with an Object-Centric Abstraction.\" arXiv preprint arXiv:2408.11816 (2024).\n\n[4] Zhou, Allan, et al. \"Policy architectures for compositional generalization in control.\" arXiv preprint arXiv:2203.05960 (2022).\n\n[5] Haramati, Dan, Tal Daniel, and Aviv Tamar. \"Entity-centric reinforcement learning for object manipulation from pixels.\" arXiv preprint arXiv:2404.01220 (2024).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "F3EmhBX9EW", "reviewer": "Reviewer_zM4S", "summary": "This paper studies object based representations for learning world models in a sample efficient manner. They propose to use SAM2 and Cutie as the frozen models which provide object vectors. These vectors along with the full frame are passed into categorical vae followed by a spatial transformer whose outputs are used to predict the next latents, rewards and termination. They show the effectiveness of their approach on atari100k, hollow knight and some continuous control tasks.", "questions": "What is the policy to collect rollouts for training the world model?\n\nHow do you decide the number of object vectors for each task?", "limitations": "", "rating": 4, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "I appreciate the idea of adding object-based inductive biases into world modelling, this allows to maintain object consistency and tracking for all objects especially smaller ones.\n\nI like that the authors opted for wide evaluation suite going beyond atari to hollow knight and continuous control.", "weaknesses": "I think the paper would benifit from comparing against unsupervised object-centric representation learning baseline such as slot attention (https://arxiv.org/abs/2006.15055) and slotformer (https://arxiv.org/abs/2210.05861). The main claim of the paper is that object based vectors help world modelling but it is not clear whether there is something special in vectors provides by SAM2/Cutie or even unsupervised methods can also help. Also world modelling has been a focus in the unsupervised object-centric representation learning community (eg. slotformer, SSWM( https://arxiv.org/abs/2402.03326)) hence it might be useful to copmare against some of these baselines.\n\nIt would be great if the authors could clarify and make the setup more clear. For example,\n- The main model figure is in the appendix, I think it should be in the main paper\n- It is not clear how the FOCUS baseline works - how do you discretize the masks from FOCUS?\n\nThe claim for discretization is to limit compounding of errors. I don't fully buy this, shouldn't errors compound in discrete autoregressive prediction too? It would be great to have a baseline which does not perform discretization so only continuous prediction to show the pros and cons\n\n\nThe authors claim that the method is sample efficient which is great, however if given access to more samples, will the method scale as well as non object-based methods?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)