Leveraging Separated World Model for Exploration in Visually Distracted Environments
Authors: Kaichen Huang
Shenghua Wan
Minghao Shao
Hai-Hang Sun
Le Gan
Shuai Feng
De-Chuan Zhan
Venue: NeurIPS 2024 poster (2024)
Keywords: unsupervised RL
separate world model
visual inputs with distractors
minimax optimization
Primary area: reinforcement_learning

Abstract: Model-based unsupervised reinforcement learning (URL) has gained prominence for reducing environment interactions and learning general skills using intrinsic rewards. However, distractors in observations can severely affect intrinsic reward estimation, leading to a biased exploration process, especially in environments with visual inputs like images or videos. To address this challenge, we propose a bi-level optimization framework named Separation-assisted eXplorer (SeeX). In the inner optimization, SeeX trains a separated world model to extract exogenous and endogenous information, minimizing uncertainty to ensure task relevance. In the outer optimization, it learns a policy on imaginary trajectories generated within the endogenous state space to maximize task-relevant uncertainty. Evaluations on multiple locomotion and manipulation tasks demonstrate SeeX's effectiveness.

Leveraging Separated World Model for Exploration in Visually Distracted Environments
Kaichen Huang 1.2 .2- Shenghua Wan 1:2,, Minghao Shao 1.2
Hai-Hang Sun 1.2 Le Gan Shuai Feng". De-Chuan Zhan School of Artificial Intelligence, Nanjing University. China 2 National Key Laboratory for Novel Software Technology. Nanjing University, China School of Cyberspace Science and Technology, Beijing Institute of Technology, China (huangkc, wansh, shaomh, sunhh} amdda.. nju..nee. edu. cn. {ganl, zaandc Jnjju. edu. en, fengshuai @bit edu.
Abstract
Model-based unsupervised reinforcement learning (URL) has gained prominence for reducing environment interactions and learning general skills using intrinsic rewards. However, distractors in observations can severely affect intrinsic reward estimation, leading biased exploration process, especially in environments with visual inputs like images or videos. To address this challenge, we propose a bi-level optimization framework named Separation-assiiet the inner optimization, SeeX trains separated world model to extract exogenous and endogenous information, minimizing uncertainty to ensure task relevance. the outer optimization, it learns policy on imaginary trajectories generated within the endogenous state space to maximize task-relevare uncertainty. Evaluations on multiple locomotion and manipulation tasks demonstrate SeeX's effectiveness.
eXplorer (SeeX).
Introduction
Unsupervised learning has rich history in com- puter vision and natural language processing, as demonstrated by methods such as 12 23 25
9. 48, 14] It leverages unlabeled, task-agnostic data to train models that can be quickly adapted to downstream tasks, addressing sample ineffi- ciency. This approach has also gained traction in unsupervised reinforcement learning (URL) 35 29, 36), where skills are developed through intrinsic motivation rather than external rewards. key advantage of URL is that learned dynam- ics models can gather prior knowledge about the environment during exploration, minimizing the need for extensive interactions during adaptation to downstream tasks.
(2,4,)
of x
W.D. scparate
separate Wecisted
Inmoiiin predictiv head
uncaraintt cximaiinn
Comparison of traditional URL meth- ods using a single world model (left) versus our separated world model (right).
Figure 1:
major challenge unsupervised reinforcement learning arises from distractors, which significantly hinder learning in complex real-world environments. For example, book-finding robot in library must identify relevant objects (like books) while ignoring irrelevant ones (such as posters or people) to effectively complete its task. These distractors exponentially expand the original state space, creating numerous redundant states that impede efficient exploration [64]. As shown in Figure Left, a single world model encodes both task-relevant and task-irrelevant information in the latent space. Even
Equal contribution. Corresponding author.
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

if the agent's state remains unchanged in 2,+1*, varying background distractors inflate uncertainty estimates. In contrast, as illustrated in ( our proposed separation world model distinctly separates task-relevant information 8+ from task-irrelevant information sT allowing for accurate uncertainty estimation unaffected by distractors.
Despite its prevalence in real-world environments, this challenge has received limited attention in unsupervised reinforcement learning. Current unsupervised RL methods: spanning data-driven, knowledge-based and competence-basee approaches (see Section 6 remain vulnerable to redundant states, which can lead agents down unproductive exploration paths 66
To address visual inputs with complex distractors, we propose a bi-level optimization framework called Separation Assisted Explorer (SeeX). We extend URLB URLB-pixels to environments with distractors, demonstrating that performance declines vary with task difficulty. Our method, SeeX, separates latent information in the world model into task-relevant and task-irrelevant components. We assume the task-irrelevant part captures transitions, using it to predict rewards and actions.
Our contributions are summarized as follows: (I) We introduce Ipparation-Assisted Explorer (SeeX), a new approach for tackling unsupervised RL tasks with complex distractors in observations. (II) We provide theoretical analysis that formalizes URL with distractors, where the policy maximizes task-relevant uncertainty while the world model minimizes environmental uncertainty. (III) We demonstrate the outstanding performance of policies learned by SeeX across various locomotion and manipulation tasks with diverse visual distractors.
Preliminary
World models are highly effective for reinforcement learning (RL) tasks with visual inputs, performing well in both simulated [20] 21 56) and robotics [67 66 environments. In unsupervised settings, agents explore using intrinsic rewards during pre- pre-training This task-agnostic data facilitates fine- tuning through a world model. URLB-pixele [50] demonstrates that unsupervised methods combined with world models outperform model-free agents. Inspired by these results, we adopted a model-based framework in our approach.
Several studies, including URLB [33] and URLB-pixels [50]. have created unified framework for various unsupervised reinforcement learning (URL) methods, advancing the field. The benchmark consists of two phases: pre-training (PT), where agents explore without task-specific rewards, and fine-tuning (FT), where they apply knowledge gained during PT with limited interactions. While uRLB URLB-pixels effectively evaluates pixel- -based performance, it does not address complex distractors in visual inputs. To fill this gap, we tested baseline methods on noisy video backgrounds and developed a novel approach to tackle this challenge.
Our focus lies in scenarios where agents must acquire skills from noisy visual observations, often disrupted by task-irrelevant videos. To model the dynamics under these conditions, we adopt the EX-BMDP framework, adaptation of Block MDP. A Block MDP can be represented
tuple M = (O, A, R,4)) where O, Z and A denotes the set of observations, latent states and actions respectively; the transition function T Z x A - A(Z) maps latent states and actions to probability distributions over next latent states; the reward function R. : O x A + [0, 1] assigns reward to iiservation-action pairs; the emission function 2/ : Z - (OO) maps latent states to probability distributions over observations. The EX-BMDP is formulated as follows:
Definition 2.1. (Exogenous Block Markov Decision Process). An EX-BMDP is BMDP such that the latent state can be decoupled into two parts 2 = (s*, " where s* E S* is endogenous state and 8 S is the exogenous state. For E Z the initial distribution and transition functions are decoupled, that is: pu()) 4+ (s+ (s ), and (='|(,,a) (st |st ( as-T
3 Bi-level Optimization for Exploration in Distracted Environments
The objective of the pre-training phase in URL is to learn world model capable of handling downstream tasks with diverse reward functions. Following previous work 13 8 52 the pre-training phase as minimax regret problem.
we formalize
Problem 3.1. (Regret Optimization for World Model) In the context of reward-free BMDP, consider world model W that defines the latent dynamics MR, which simulates the real dynamics M
2

---

under a reward function R. Let (=)(T) denotes the expected value of policy under dynamics M with reward function R. The regret of policy π is defined as R M VV= (M*(a") Vaa (x), x² arg maxx' M (Maa)"), (T'), which measures the performance gup between x and the optimal policy. The optimization goal to find the world model that robust to different possible reward functions, minimizing the regret of world model policy under real dynamics:
min max REGRET MMMTTR..... M, = arg max V MR (T( R
(I)
In the absence of the reward function during pre-training, computing regret under specific reward functions becomes infeasible. To address this challenge, we propose transforming regret into novel,
reward-free objective using Simulation Lemma 28 52]. Our proposition is outlined below.
Denote the learned latent dynamics in the world model as T and the true latent dynamics as T. For any reward function R, the regret of the optimal world model is bounded by:
Proposition 3.2.
27
REGRET ( MM)
S
2,andla
(1 11-79 7)2
[V (((
+ E 2.a
As outlined in Proposition 3.2 the optimal world model policy exhibits low regret if the latent state-action distribution of both गे and R in M is accurately captured by T. However, during unsupervised pre-training, the reward function remains inaccessible, preventing from directly obtaining these distributions. This is because the state-action distribution induced by Th and # is inherently dependent the reward function. To address this challenge, introduce exploration policy <xl,, specifically designed to maximize the expected error (in terms of total variation distance) of the latent dynamics model:
expl arg max E, [and(m.MM)
(TT
This enables to derive upper bound the regret that is independent of the reward function:
Rf M
4y
E s.a
(T(F7122,0, for all Rt
(1 7)2 2.a
In the context of EX-BMDPss the transition dynamics can be decomposed into two components: T(z'l 2, a) = (s+ ast.a)T_ (s where (st, sT), st represents the endogenous state, and s represents the exogenous state. To apply EX-BMDP visually distracted environments, we further assume that the reward depends solely s*, implying that decisions made based only endogenous state.
To model this separation, we construct the world model as two distinet components: the endoge- part M and the exogenous M-, corresponding the task-relevant and task-irrelevant processes, respectively. Their associated latent dynamics are denoted by T. and T_. Utilizing this decomposition, we can derive lower bound for the policy optimization target:
TV (T (7)11)
E 22
TV (110 ()167(1)77
(.187), (((
(5)
sl is aa M
TV (f((staa)) T+)sst) .a)]
> aa
The strategy's optimization objective a form of model uncertainty. Estimating model uncertainty for the URL serves reward for the exploration policy, which is crucial for learning. However, under the EX-BMDP assumption, directly maximizing the reward the overall model error introduces significant bias. This is because the objective models not only the endogenous part but also the

---

exogenous part, while the strategy is solely based on s* leading to inefficient exploration. For the model optimization target, using triangle inequality, we have:
E 2,0rd(
TV (2)
TV [T+ "atl:t..),,7)411)) .))]
M M.
(6)
+ E -(
TV (
With these bounds, we can reformulate the problem as a bi-level optimization problem:
Problem 3.3. (Bi-level Optimization of World Model Error) Consider reward-free EX-BMDP with latent dynamics functions T=,, T- and an exploration policy x, assuming the reward depends only on the endogenous state. Let T T+ x T- denote the true latent dynamics. We use a bi-level optimization: the inner optimization finds the world model minimizing the error, while the outer optimization maximizes the endogenous error with respect to the exploration policy.
Inner:
min E
TV
(1) .a)]
(-1st, T+(sss
a
M,
Ti
+ min E 7_
TV (,
Outer: max E, ,a-d(
TV (11st (111 a))]
Problem 3.3 optimizes the world model error surrogate for Problem 3.1 seeks world model with low prediction error for both endogenous and exogenous components, guided by an exploration policy that maximizes endogenous error. This ensures the optimal policy generalizes well across varying exogenous distributions, regardless of the future reward function. Next, we introduce practical framework solve Problem 3.3
Practical Implementation of SeeX
4.1 Control with Separation-assisted Latent Dynamics
Separated World Model. For high- digh-dimensional inputs like images and videos, model-based frameworks encode past experiences into latent representations to predict future sequences [59 65 19 Building on this, we propose a separated world model to disentangle task-relevant (endogenous) and task-irrelevant (exogenous) information. As shown in Figure 2 the endo-encoder h Fig+ (o:) and exo-encoder h, = Eg (o:) extract these components. Two transition models handle dynamics: qe- (s+ | 1. ((-1) for task-relevant states and qo (57 { for task-irrelevant ones. Additionally, inference models po+ (s |s hx and pg- (s, 1. h provide posterior estimates.
use an auxiliary reconstruction loss to refine the observation encoder, ensuring captures sufficient information from o, by maximizing the mutual information I(o;; 2.). The IM algorithm 6 provides the Barber-Agakov lower bound:
Observation Model. Some model-based approaches 20 49]
I(o =)) 2 Exlo...)) (n (o. (o + where H((((,)) is the entropy of the observation distri- bution. In our separation setting, this bound becomes I(o, st 87 E pllor.e .e.. IIn (o, 8 )}-
Inspired by 163 18 24]. We design observation model go no(olsst, , (o,t .m, ) ~ Dø (s an exo-decoder (o,,m,) (g- and a mask-mixing model m, = Me(m," .m))) of and o, are expected to reconstruct the task-relevant and task- iusk-irrelevant parts of the original observation, while m and m, are corresponding masks. So (o(o) (o can be implemented as ô, = my 0 + +(1 +11my @ Considering insights from [18 64) that task-relevant information occupies only small portion of the observation, design exogenous reconstruction (Exo-Rec) model 0r50 go(o, to ensure that s contains the vast majority of information.
endo-decoder
World Model Optimization. Similar to a variational autoencoder (VAE) 30 51), all model components are trained jointly by maximizing the evidence lower bound (ELBO), Our optimization

---

exogenous branch
imaginary trajectories
tmttr Vur Warthhiskeeeend
6
endogenous branch
(a) inner optimization of world model
(b) outer optimization pollcy
Figure 2: (a) Our separated world model comprises task-relevant and task-irrelevant branches, with ensemble predictive heads providing reward signals. (b) Imaginary trajectories within the endogenous branch enhance sample efficiency by reducing real-world interactions.
objective extends that of Dreamer 20
L = Epp (4.66 CR + fk,)). The detailed formulations of each term are presented follows, and the derivations are shown in Appendix B.1
co = In (orls) s.) go
(8)
C = In (r(r)(s²)
(9)
Ckt KL(poi (s Ilaa (st 11 + KI,po (s5 18 1loo (#, 18 at-1))
Policy Optimization. To enhance sample efficiency, as shown in Figure 2 (b), we apply the
policy optimization strategy of [20] to learn a parametric policy with imaginary trajectories. More specifically, we design an actor ( (s and a reward model go (era/ss (rrss")) The reward model fits the true reward function in the FT phase. The actor selects action based the current endogenous state
s7 under the guidance of intrinsic reward in the PT phase or the reward model in the FT phase.
Intrinsic Reward: Compute Endogenous Uncertainty under Separation View
Crafting an effective intrinsic reward function is paramount during the pretraining (PT) phase. A well- designed intrinsic reward can steer the agent towards exploring states with the highest uncertainty, maximizing skill acquisition. A common practice 46 47 [4] is to maximize the mutual information : $ c. where & represents the true unknown dynamics. This objective is motivated by the fact that mutual information quantifies how comprchensively the trajectory explores the environment.
Proposition 4.1. (Single step's 's mutual information) If only we find the policy Texpl that maximize every single step's 's mutual information Et- I1 at), then the whole trajectories's 's mutual information I( s Teppl) will be also maximized. Proof in Appendix B.2
With Proposition the optimization objective can be written as:
Texpl arg max Test
I(htari {ls,
(10)
arg max Teyl
H(h lst tasttaa at)
Σ
H( = ,. (t))
(11)
t=
t=
To approximate the unknown ç. we design a set of predictive heads {yu (Tott Kek,, of a 1.k which are implemented conditional Gaussians ((pEE = Eu, aa)), o²). the conditional entropy does not depend on state or action [54). Suppose that p(=) is an uniform distri-
Given fixed variance,
k= H(h E Ek. st a.). Since the bution on T then we have H Hh |st at) * ΣK
5

---

Stand
Walker Walk
Walker Run
Walker Flip
dormalize Return
1.00
1.0
0.6
0.50
0.50
0.25
0.25
0.0
0.0
0.00
0k
50k
100k
0k
50k
100k
0k
50k
100k
0k
50k
100k
Quadruped Stand
Quadruped Walk
Quadruped Run
Quadruped Jump
Normalized Retum
1.0
0.6
0.75
0.4
0.50
0.5
0.2
0.2
0.25
0.0
0.0
0,00
0k
50k
100k
0k
50k 100k
0k
50k
100k
0k
50k
100k
Plan2Explore
Figure 3: We show fine-tuning (FT) performance curves of SeeX and baselines across two domains and eight tasks. Pre-training (PT) used 2M frames, FT 100K. Normalized returns are benchmarked against the expert baseline 33 with mean (solid line) and variance (shaded area).
marginal entropy in Equation [11 lacks a closed-form expression amenable optimization, em- ploy the empirical variance across ensemble means as substitution. is also an approximation to the world model TV error Problem 3.3 We define the intrinsic reward as rf Var( Var((htt,. )e..:))
and employ the exploratory policy Texpl that maximizes ritl approximation of Tap in Equa-
tion 10
5 Experiments
All experiments are conducted with at least three seeds and evaluated for 10 episodes. We conduct experiments to answer the following questions: (1) Can SeeX outperform other counterparts in URLB? (2) How do pre-training steps affect final performance? (3) Can SeeX give a reasonable model uncertainty estimation? (4) How do different components affect the model training? (5) Can the pre-trained world model and policy generalize to OOD distractors? (6) Can data augmentation handle visual distractors?
Environments. Following the URLB evaluation, we select three domains: Walker, Quadruped, and Jaco Arm, spanning twelve downstream tasks: Walker (stand, walk, run, flip), Quadruped (stand, walk, run, jump), and Jaco (reach-top-right, reach-top-left, reach-bottom-right. reach-bottom--let)) These tasks vary in difficulty, providing well-rounded performance assessment. The Jaco Arm domain is particularly challenging due to its multi-joint structure and sparse rewards, only granted upon catching the red ball. Visualizations of the environments are provided Appendix Agents receive only visual inputs (64, 64, 3), with episodes lasting 1000 frames and an action repeat of R Pre-training utilizes up to 2M frames, and fine-tuning employs 100K frames, consistent with URLB-pixels.
Datasets. Leveraging the large-scale and high-quality Kinetics dataset 227 we construct two sub- datasets for our experiments: the Driving-car dataset, commonly used in previous works [18] 63], and the Random-video dataset, composed videos randomly selected from other Kinetics classes. To evaluate generalization capabilities, we conducted experiments on the Random- video dataset Section 5.5 (Note that we only use the Candom-vide here). This zero-shot transfer assessment gauges the effectiveness of the policy on unseen random video scenarios.
Baselines. We choose seven different unsupervised methods as baselines, which can be categorized into three types: Knowledge-buuk based: ICM [46], RND [10), LBS [4] and Plan2Explore [54] maximize prediction error to better understand the world; Data-basedl: APT (9) encourages exploration by maximizing entropy: Competence-based. DIAYN [16] and APS 38 maximize mutual information to achieve diverse discovery and generalization. The implementation detail and more concrete introduction of the above methods shown in Appendix D.2

---

Walker Stand
Walk
Algorithm
Algorithm Y
Michad
20K
50K
100K
200K
100K
400K
SoeX
APS
APT
APS
SeeX
0.29
0.95
0.96
DWAYN
APT
0.37
0.43
0.60
0.71
0.81
0.91
1,30
0.99
0.95
ICM
DIAYN
0.40
0.41
0.48
0.40
0.65
0.53
0.98
0.36
0.51
LBS
ICM
0.33
RND
LBS
0:44
0.40
0.58
0.87
0.58
0.96
0.58
PlanExxplore
RND
0.26
0.28
0.41
0.76
0.81
0.96
0.96
0.98
0.99
0.0 0.2 0.4 0.6 0.8
PlaniTagiore
0.50
0.46
0.54
0.69
0.86
0.58
0.93
1.00
0.58
P(X > Y)
(a) Jaco Performance
(b) Walker-Stand/Wall performance vs. pretraining steps
Figure 4: (a) Average performance of four Jaco arm tasks. Each row shows comparative probabilities with 95% confidence intervals, indicating Algorithm X outperforms Algorithm [. Probabilities are based on 150 runs (50 per seed across 3 seeds) per task for robust evaluation. (b) Performance varies with pretraining steps. Given the small gap at 500k on Walker-Stand a table highlights details, with red marking over 90% expert performance.
5.1 Evaluation on URLB with Distractors
The performance curves in Figure 3 show the normalized fine-tuning performance of SeeX and baseline methods pre-trained on 2 million frames. SeeX consistently outperforms or matches other methods across tasks, achieving expert-level performance in the walker-wal walk task. Notably, SeeX displays higher mean performance and greater variance, reflecting the positive correlation typical in reinforcement learning. In the quadruped domain, even SeeX's s lowest performance exceeds the highest curves of other methods, trend seen across quadruped tasks. Plan2Explore shows relatively weaker results, especially the walker-stand and walker-walk tasks, likely due to relying solely on predictive heads based the entire latent space. In contrast, our method demonstrates significant performance improvements. Other baselines like RND and APT perform reasonably well in simpler tasks but struggle with more complex tasks like walker-run and walker-flip. This indicates that using prediction errors, entropy maximization, or mutual information alone is inadequate. In contrast, our separated world model effectively extracts relevant information to enhance policy training.
5.2 Can SeeX Give a Reasonable Model Uncertainty Estimation?
As stated in Section 4.2 model uncertainty is an approximation to the world model TV error in Prob- lem D Thus we evaluate the ac- curacy of dynamics fitting for SeeX and Plan2Explore (both incorporating model uncertainty) with pre-trained models of 1M frames. To ensure a comprehensive assessment, we em- ploy two policies: a random policy and SeeX's exploration policy (IM frames). As illustrated in Table [ con- sistently across policies and domains, U(++) values are lower than (2) val- ues. This indicates that SeeX achieves a more accurate estimation of true dy- namies compared Plan2Explore, potentially contributing to its superior performance relative to other baseline methods. Furthermore, notable observation is that U(s-) values significantly exceed U(++) values. This suggests that the total environmental uncertainty remains unchanged; rather, most uncertainty is concentrated in go- and Pa- due to SeeX's 's effective capture of task- iusk-irrelevant distractor transitions.
Policy
Method
Walker
Quadruped
Jaco
U(s*)
0.04:1000
0.17 0.00
0.131:0.00
random
U(s U(=)
0.64
33.0#11.0
0,48+00.3
0.31-10.01
1.49:10.09
0.40-10.01
U(s*)
0.05=0000
0.18:10.00
0.14
SeeX
U(s
0.65+00.03 45.0++50.0
0.52+0..3
U(2)
0.35++0.01
1.78
0.44
Table 1: Uncertainty estimation of SeeX and Plan2Explore on the Driving-car dataset. Using two policies, we collected 1000 distinct states and reported the mean and standard de-
viation of uncertainty. U(s)), ), U(s"), ). and (z) denote the uncertainty for SeeX (endogenous and exogenous states) and Plan2Explore (latent belief state), respectively.
5.3 How do pre-training steps affect final performance?
Building upon the promising results from the previous section, where SeeX outperformed seven baseline methods and achieved expert-level performance from 3 (trained in a distractor-fres environment, representing the upper bound for our setting) in some tasks, we delve into intriguing
7

---

Complete Model
Ablated Model
o-walker
a-quad
a-jaco
Enviromment
Normalized Return
0.8
0.10
Quadruped
0:49
0.15
0.36
0.37
0.47
0.08
Walker
0.80
0.65
0.69
0.82
0.6
(a) Ablation results of different modules.
0.06
0.4
0.04
Walker
0.2
0.1
0.02
015
015
0.0
0.0
0.00
ge0
deg
deg
ded
del
606
140 1ie
J400
(b) Ablation of different policy design.
(c) Ablation on (weight of Exo-Rec term).
Figure Ablation results of SeeX with 500k fixed pretraining frames: (a) separation design. Exo-Rec term, and predictive head values; (b) different policy design: T(s+) (SeeX) and (( (SeeX both-branch); (c) impact different (Exo-Rec weight) across three domains. Due Jaco's task complexity, we ran 50 trials per seed and reported the top 30 mean returns.
question: how does the fine-tume fine-tuning performance SeeX vary with different pre-training step sizes? More specifically, can we achieve fine-tuner tunettundd performance with as few pre-training steps possible achieve more than 90% of expert performance? We conduct experiments on walker-stand and walk. For all methods, we assessed the fine-tuning performance using pre-training frames of (100k, 500k, Im, 2m). To evaluate SeeX on more granular level, we add choices (20k, 50k, 200k, 300k, 400k). As shown in Figure performance (red dotted line) within 500k frames. contrast, no baseline achieves this level, even with 2M pre-training frames. For the simpler Walker-Stand task. table offers detailed analysis. SeeX requires only 100k frames of random exploration to reach 90% expert performance, while other baselines need significantly more pretraining.
for the Walker-Wail task, SeeX reaches 90% expert
5.4 How do different components affect the model training?
Ablation of modules. We first examine the impact of various modules in SeeX, including the separate world model, Exo-Rec design, and the number of predictive heads (K). As shown in Figure (a). removing either the separate world model Exo-Rec significantly reduces performance, underscoring the effectiveness of these designs. A small value of K hinders accuracy, while performance improves with more heads, confirming our intuition. To balance estimation accuracy and computational cost, we choose K = 5.
Ablation of policy. In SeeX, use only the endogenous state s+ to predict actions. To validate this design, we compare two policies: T(8*) (SeeX) and 1 = (SeeX-both-branch As shown in Figure (b)),, SeeX consistently outperforms -exX-both-branchh particularly in the last three challenging tasks. This demonstrates that relying solely on endogenous information is beneficial in environments with moving distractors.
Ablation of a. We conducted experiments to assess the impact of Exo-Rec's weight a. As shown in 5) the optimal value of a varies across domains. For Jaco and quadruped tasks, smaller (e.g.. I) enhances performance, whereas walker tasks benefit from a larger weight (e.g., 2 or 3) to effectively extract exogenous information into
5.5 Can the Pre-trained World Model and Policy Generalize to OOD Distractors?
In this subsection, evaluate the effectiveness of SeeX in addressing the "generalize to distractors from other distributions (OOD distractors)" challenge. To assess generalization ability, we compare SeeX with other baselines on four different walker tasks and report the average normalized return. We pre-trained and fine-tune tuned agents on the driving-car dataset and evaluated their performance on both driving-car and random-video datasets. The reported results represent the average normalized return across four walker tasks, with 50 runs conducted for each seed. As shown in Figure (b), SeeX outperforms baselines on test distractor datasets and shows minimal performance drop under distribution shifts, highlighting its strong generalization.

---

Method
Design
Stand
Flip
Method
Driving-car
Random-video
Drop(()
0.99
0.49
0.78
SeeX
0.82
0.75:: 0.07
-0.09
SeeX
DA
1.00
0.96
0.53
0.81
APS
0.55
0.27# 0.06
-0.51
pDA
0.97
APT
0.64
0,55: 0.07
-0.14
PDA-IDA
0.97
0.92
0.36
0.70
DIAYN
0.63
0.34: 0.08
-0.46
0.97
0.77
0.29
0.56
ICM
0.69
0.56+ 0.07
Plan2Explore
PDA
fDA
0.95
0.98
0.49
0.23
0.50
Plan2Explore
LBS
0.70
0.69
0.58+ 0.06
0.61 0.08
-0.13
-0.16
PDA+IDA
0.95
0.56
0.24
0.53
RND
0.64
0.49+ 0.06
-0.23
(a) Impact of data augmentation techniques.
(b) Generalization ability to OOD distractors.
Figure 6: (a) The impact of DA performance in PT (pDA) and FT (IDA) stages. (b) Generalization ability to distractions from other distributions. The rightmost column indicates the percentage drop in performance caused by the distribution shift.
5.6 Can data augmentation handle visual distractors?
70 explore data augmentation (DA) RL as an effective strategy for visual generalization. This subsection highlights the differences between DA and our approach. Our bi-level separation framework extracts task-relevant information, similar how DA captures task-relevant representations. However, since few approaches use DA for exploration, we integrated it into Plan2Explore and SeeX to assess its impact on performance in the moving distractor setting. Specifically, we applied the classic random shift augmentation from [73 (4-pixel padding) four times per image. To evaluate DA's effectiveness, we tested it during both pretraining (pDA) and finetuning (fDA) on the walker task. The results in Figure (a) reveal: (1) fDA improves performance by mitigating distractors. (2) SeeX's separation design outperforms fDA on tasks with moving distractors. (3) pDA reduces performance, likely by disrupting world model learning, further investigation is planned.
Numerous studies [
6 Related Work
Learning a dynamics model of the environment is promising way to tackle the problem of low sample efficiency, and has achieved impressive results in various tasks including continuous control as well as discrete control. To ensure the fairness of comparison, we adopted the official version of Plan2Explore and the implementation baseline methods into unified Dreamer-like like framework. Plan2Explore makes framework and takes the uncertainty of prediction of the next latent state as the intrinsic reward. In work, focus on the scenarios with complex visual distractors and design an intuitive separated world model to capture the exogenous and endogenous states respectively, and learn policy with imaginary trajectories in latent space.
Model-based Control.
URLB-piss Uii.Bl, which combines all
of the dreamerv2
Unsupervised RL. In recent years, there are many works that tried promote the unsupervised
representation learning manner in various fields including computer vision (CV) 12 23 25 and Additionally, there have been numerous works These works encourage the 72 26), however, 58
natural language processing (NLP) 9 48 14)
applying RL algorithms to real-world applications, such as 37 RL community explore the more efficient way of learning 34
68
44
32
53
these works still need to optimize an extrinsic reward. Recently, there have been works to adopt a pure unsupervised manner (reward-free pre-training followed by reward-specific fine-tuning), which can be categoried into three kinds. (i) Data-based: maximal entropy RI has enabled agent for
diverse exploration and data 39
55
7I:: (ii) Knowledge-bass -base:: increase knowledge about the
47); (iii) competence-based methods based on mutual information [16, 22 38 55] shows capability for diverse discovery and generalization. URLB- pixels offers a unified basee--basd framework to implement some methods referred to above, we make use of this benchmark to serve the baseline. shared limitation of existing unsupervised exploration methods lies in their exclusive reliance on state information for exploration, lacking the ability to explicitly extract endogenous information. This inherent limitation renders these methods susceptible to performance degradation when confronted with visually rich environments containing complex distractors. Different from the above methods, our method designs an intuitive separated model to separate exogenous and endogenous information and only collect imaginary trajectories in
world with self-supervised prediction 46

---

endogenous latent space to train policy. As result, we largely improved performance compared to the other methods.
Learning with Noisy Observations. To address the issue of learning with noisy observations, diverse approaches have emerged, broadly categorized into four main strategies: (1) Use data
augmentation methods to mitigate exogenous noises [72 17 74 77: (2) Learning task-relevant representations with bisimulation metrics 75
40] (3) Design auxiliary tasks to extract endogenous information 4 15 31]: (4) Take actions rewards as discriminating factors to separate exogenous Our proposed method, SeeX, falls into the intersection of the last two categories, combining the benefits of auxiliary task design and discriminating factors to effectively handle noisy observations complex real-world environments.
and endogenous information 18, 64 45 41
Conclusion
In this paper, consider the problem of learning a policy in the visual unsupervised reinforcement learning (URL) setting with moving distractors. To tackle this intricate scenario, we introduce SeeX, a bi-level optimization framework that leverages separated world model and task-relevant uncertainty maximization to mitigate the impact of distractors and enhance exploration efficiency, Further- more, Our theoretical analysis formalizes URL with distractors: the policy maximizes task-relevant uncertainty to drive exploration, while the world model minimizes environmental uncertainty to reduce distractor influence. SeeX utilizes separated world representation disentangle exogenous and endogenous factors from the original observation domain. Policy training is then conducted exclusively on imaginary trajectories generated within the endogenous latent representation. Exten- sive experiments on the DMC-suite benchmark demonstrate that SeeX outperforms other baseline methods. Ablation studies provide insights into the contributions of individual components to our final performance.
Limitations and Future Work. Our work presents several areas for improvement: (I) Our work focuses on DMC tasks, leaving real-world applications like self-driving and navigation challenges for future exploration. (II) Distractors fall into four types based on their impact on rewards and actions: task-irrelevant action-independent (our focus), task-relevant + action-indepent indepenenn,,,, task-irrelevant + action- detton-dependent, and task-relevant + action-dependent. The latter three will be addressed in future work. (III) Using only s+ for policy learning is beneficial in our setting. However, for the task-relevant action-indeen intioin-nndeeddeent case (e.g., multi-agent systems), we propose incorporating some into policy optimization as potential solution.
Acknowledgments and Disclosure of Funding
This work was partially supported by the National Science and Technology Major Project under Grant No. 2022ZD01 14805, Collaborative Innovation Center of Novel Software Technology and Industrialization, NSFC (623761 18, 62006112. 12, 62250069, 61921006), and the Postgraduate Research & Practice Innovation Program of Jiangsu Province.
10

---

References
[1] Rishabh Agarwal, Max Schwarzer, Pablo Samuel Castro, Aaron C Courville, and Mare Belle- mare. Deep reinforcement learning at the edge of the statistical precipice. Advances in neural information processing systems, 34:29304 24293004-29320, 2021.
[2] Alexander A Alemi, lan Fischer, Joshua V Dillon, and Kevin Murphy. Deep variational information bottleneck. arXiv preprint arXiv: 16/2.006112004000 2016.
[3] Marcin Andrychowicz, Filip Wolski, Alex Ray, Jonas Schneider, Rachel Fong, Peter Welinder, Bob McGrew, Josh Tobin, OpenAl Pieter Abheel, and Wojciech Zaremba. Hindsight experience replay. Advances in neural information processing systems, 30, 2017.
[4] Adrià Puigdomènech Badia, Bilal Piot, Steven Kapturowski, Pablo Sprechmann, Alex Vitvit- skyi, Zhaohan Daniel Guo, and Charles Blundell. Agent57: Outperforming the atari human benchmark. In International conference machine learning, pages 507-517. PMLR, 2020.
[5] Bowen Baker, 1lge Akkaya, Peter Zhokov, Joost Huizinga, Jie Tang, Adrien Ecoffet, Brandon Houghton, Raul Sampedro, and Jeff Clune. Video pretraining (vpt): Learning to act by watching unlabeled online videos. Advances in Neural Information Processing Systems, 35:24639-24444 -46544, 2022.
[6] David Barber and Felix Agakov. The im algorithm: variational approach information maximization. Advances in neural information processing systems, 16(320):201, 2004.
[7] David Bertoin, Adil Zouitine, Mehdi Zouitine, and Emmanuel Rachelson. Look where you look! saliency-guided q-networks for visual rl tasks. In Thirty-sixth Conference on Neural Information Processing Systems (NeurIPS 2022), 2022.
[8] Michael Beukman, Samuel Coward, Michael Matthews, Mattie Fellows, Minqi Jiang, Michael Dennis, and Jakob N. Foerster. Refining minimax regret for unsupervised environment design. CoRR, 0tbs/2402.12284, 2024.
[9] Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models few-shot learners. Advances in neural information processing systems, 33:1877-1901, 2020.
[10] Yuri Burda. Harrison Edwards, Amos Storkey, and Oleg Klimov. Exploration by random network distillation. arXiv preprint arXiv:1810.128%) 2018.
[11] Lili Chen, Kevin Lu, Aravind Rajeswaran, Kimin Lee, Aditya Grover, Misha Laskin, Pieter Abbeel, Aravind Srinivas, and Igor Mordatch. Decision transformer: Reinforcement learning via sequence modeling, Advances in neural information processing systems, 34:1 50844110777 2021.
[12] Ting Chen, Simon Kornblith, Mohammad Norouzi, and Geoffrey Hinton. A simple framework for contrastive learning of visual representations. In International conference machine learning, pages 1597-1607. PMLR, 2020.
[13] Michael Dennis, Natasha Jaques, Eugene Vinitsky, Alexandre M. Bayen, Stuart Russell, Andrew Critch, and Sergey Levine. Emergent complexity and zero-shot transfer via unsupervised environment design. In Hugo Larochelle, Marc' Aurelio Ranzato, Raia Hadsell, Maria-Florina Balcan, and Hsuan-Tiee Lin, editors, Advances Neural Information Processing Systems 33: Annual Conference on Neural Information Processing Systems 2020, NeurlPS 2020, December 6-12, 2020, virtual, 2020.
[14] Jacob Devlin, Ming-Wor Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint arXiv: 18/0. 040055 2018.
[15] Yonathan Efroni, Dylan J Foster, Dipendra Misra, Akshay Krishnamurthy, and John Lang- ford. Sample-efficient reinforcement learning in the presence of exogenous information. In Conference Learning Theory, pages 5062-5127. PMLR, 2022. 2022.

---

[16] Benjamin Eysenbach, Abhishek Gupta, Julian Ibarz, and Sergey Levine. Diversity is all you need: Learning skills without a reward function. arXiv preprint arXiv:1802.060700 2018.
[17] Jiameng Fan and Wenchao Li. Dribo: Robust deep reinforcement learning via multi- information bottleneck. In International Conference on Machine Learning, pages PMLR, 2022.
[18] Xiang Fu, Ge Yang. Pulkit Agrawal, and Tommi Jaakkola. Learning task informed abstractions. In International Conference on Machine Learning, pages 3480-3491. PMLR, 2021.
[19] David Ha and Jürgen Schmidhuber. World models. arXiv preprint arXiv: 10122 2018.
[20] Danijar Hafner, Timothy Lillicrap. Jimmy Ba, and Mohammad Norouzi. Dream to control:
Learning behaviors by latent imagination. arXiv preprint 191. 2019.
[21] Danijar Hafner, Timothy Lillicrap, Mohammad Norouzi, and Jimmy Ba. Mastering atari with discrete world models. arXiv preprint arXiv:2010.02193. 2020.
[22] Steven Hansen, Will Dabney, Andre Barreto, Tom Van de Wiele, David Warde-Farley, and Volodymyr Mnih. Fast task inference with variational intrinsic successor features. arXiv preprint arXiv: /100//90600.0 05030, 2019.
[23] Kaiming He, Haoqi Fan, Yuxin Wu, Saining Xie, and Ross Girshick. Momentum contrast for unsupervised visual representation learning. Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pages 9729-97338 9738, 2020.
[24] Kaiming He, Georgia Gkioxari, Piotr Dollár, and Ross Girshick. Mask r-cnn. In Proceedings of the IEEE international conference on computer vision, pages 2961-2969 2017.
[25] Olivier Henaff. Data-efficient image recognition with contrastive predictive coding. In Interna- tional conference on machine learning, pages 4182-4192, PMLR, 2020.
[26] Michael Janner, Qiyang Li, and Sergey Levine. Offline reinforcement learning as one big sequence modeling problem. Advances in neural information processing systems, 34:1273- 1286, 2021.
[27] Longlong Jing and Yingli Tian. Self-supervised visual feature learning with deep neural net- works: A survey. IEEE transactions on pattern analysis and machine intelligence, 43(11):4037- 4058, 2020.
[28] Michael J. Kearns and Satinder Singh. Near-optimal reinforcement learning in polynomial time. Mach. Learn., 49(2-3):209 232, 2002.
[29] Seongun Kim, Kyowoon Lee, and Jaesik Choi. Variational curriculum reinforcement learning for unsupervised discovery of skills. arXiv preprint arXiv:2310. /9424, 2023.
[30] Diederik P Kingma and Max Welling. Auto-encoding variational bayes. arXiv preprint arXiv:1312.6114, 2013.
[31] Alex Lamb, Riashat Islam, Yonathan Efroni, Aniket Didolkar, Dipendra Misra, Dylan Fos- ter, Lekan Molu, Rajan Chari, Akshay Krishnamurthy, and John Langford. Guaranteed dis- covery of controlennogenen latent states with multi-step inverse models. arXiv preprint arXiv:2207.08229, 2022.
[32] Michael Laskin, Aravind Srinivas, and Pieter Abbeel. Curl: Contrastive unsupervised repre- sentations for reinforcement learning. In International conference on machine learning, pages 5639-5650. PMLR, 2020.
[33] Michael Laskin, Denis Yarats, Hao Liu, Kimin Lee, Albert Zhan, Kevin Lu, Catherine Cang, Lerrel Pinto, and Pieter Abbeel. Urlb: Unsupervised reinforcement learning benchmark. arXiv preprint arXiv:2110.15191, 2021.
[34] Misha Laskin, Kimin Lee, Adam Stooke, Lerrel Pinto, Pieter Abbeel, and Aravind Srinivas. Reinforcement learning with augmented data. Advances neural information processing systems, 33: 33:19884-19895, 2020.
12

---

[35] Hojoon Lee, Koanho Lee, Dongyoon Hwang, Hyunho Lee, Byungkum Lee, and Jaegul Choo. On the importance of feature decorrelation for unsupervised representation learning in reinforcement learning. In International Conference on Machine Learning, pages 18988-19009. PMLR, 2023.
[36] Juncheng Li, Xin Wang. Siliang Tang, Haizhou Shi, Fei Wu, Yueting Zhuang, and William Yang Wang. Unsupervised reinforcement learning of transferable meta-skills for embodied navigation. In Proceedings of the IEEE/CVF Conference on Computer Vision and Puttern Recognition, pages 12123-12132, 2020.
[37] Shuang Li, Yanghui Yan, Ju Ren, Yuezhi Zhou, and Yaoxue Zhang. A sample-efficiee actor-critic algorithm for recommendation diversification. Chinese Journal of Electronics, 29(1):89-999 96, 2020.
[38] Hao Liu and Pieter Abbeel. Aps: Active pretraining with successor features. In International Conference on Machine Learning, pages 6736-6747. PMLR, 2021.
[39] Hao Liu and Pieter Abbeel. Behavior from the void: Unsupervised active pre-training. Advances in Neural Information Processing Systems, 34: 144:1889-8888 18473, 2021.
[40] Qiyuan Liu. Qi Zhou, Rui Yang, and Jie Wang. Robust representation learning by clustering with bisimulation metries for visual reinforcement learning with distractions. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 37, pages
2023.
[41] Yuren Liu, Biwei Huang, Zhengmao Zhu, Honglong Tian, Mingming Gong, Yang Yu, and Kun Zhang. Learning world models with identifiable factorization. Advances Neural Information Processing Systems, 36, 2024.
[42] Guozheng Ma, Zhen Wang, Zhecheng Yuan, Xueqian Wang, Bo Yuan, and Dacheng Tao. A comprehensive survey of data augmentation in visual reinforcement learning. arXiv preprint arXiv:2210.045611 2022.
[43] Pietro Mazzaglia, Ozan Catal, Tim Verbelen, and Bart Dhoedt. Curiosity- driven exploration via latent bayesian surprise. In Proceedings of the AAAI conference on artificial intelligence, volume 36, pages 7752-7760, 2022.
[44] Haowei Meng, Ning Xin, Hao Qin, and Di Zhao. A recursive dri-based resource allocation method for multibeam satellite communication systems. Chinese Journal of Electronics, 1365):1286-1295, 2024.
[45] Minting Pan, Xiangming Zhu, Yunbo Wang, and Xiaokang Yang. Iso-dream: Isolating and leveraging noncontrollable visual dynamics in world models. Advances in neural information processing systems, 15:23178-2319) 2022.
[46] Deepak Pathak, Pulkit Agrawal, Alexei A Efros, and Trevor Darrell. Curiosity-driver exploration by self-supervised prediction. In International conference machine learning, pages 2778- 2787. PMLR, 2017.
[47] Deepak Pathak, Dhiraj Gandhi, and Abhinav Gupta. Self-supervised exploration via disagree- ment. In International conference machine learning. pages 5062-501 1071. PMLR, 2019.
[48] Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, al. Language models are unsupervised multitask learners. OpenAl blog, 188)99 2019.
[49] Rafael Rafailov, Tianhe Yu, Aravind Rajeswaran, and Chelsea Finn. Visual adversarial imitation learning using variational models. Advances in Neural Information Processing Systems, 34:3016- 3028, 2021.
[50] Sai Rajeswar, Pietro Mazzaglia, Tim Verbelen, Alexandre Piché. Bart Dhoedt, Aaron Courville, and Alexandre Lacoste. Mastering the unsupervised reinforcement learning benchmark from pixels. In International Conference on Machine Learning, pages
PMLR, 2023.
[51] Danilo Jimenez Rezende, Shakir Mohamed, and Daan Wierstra. Stochastic backpropagation and approximate inference deep generative models. In International conference on machine learning, pages 1278-1286. PMLR, 2014.
13

---

[52] Mare Rigter, Minqi Jiang, and Ingmar Posner. Reward -free curricula for training robust world models. CoRR, abs/2306.09205. 2023.
[53] Max Schwarzer, Ankesh Anand, Rishab Goel, R Devon Hjelm, Aaron Courville, and Philip Bachman. Data-efficient reinforcement learning with self-predictive representations. arXiv preprini arXiv:2007.05992, 2020.
[54] Ramanan Sekar, Oleh Rybkin, Kostas Daniiiddis, Pieter Abbeel, Danijar Hafner, and Deepak Pathak. Planning to explore via self-supervised world models. In International conference on machine learning, pages 8583-8592, PMLR, 2020.
[55] Younggyo Seo, Lili Chen, Jinwoo Shin, Honglak Lee, Pieter Abbeel, and Kimin Lee, State en- tropy maximization with random encoders for efficient exploration. In International Conference on Machine Learning, pages 9443-94544 PMLR, 2021.
[56] 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, pages 1332-1344. PMLR, 2023.
[57] Archit Sharma, Shixiang Gu, Sergey Levine, Vikash Kumar, and Karol Hausman. Dynamics- aware unsupervised discovery of skills. arXiv preprint orXiv:1907.0110 2019.
[58] Adam Stooke, Kimin Lee, Pieter Abbeel, and Michael Laskin. Decoupling representation learning from reinforcement learning. International Conference on Machine Learning, pages 1870-9879. PMLR, 2021.
[59] Richard S Sutton. Dyna, integrated architecture for learning, planning, and reacting. ACM Sigart Bulletin, 164):160-163, 1991.
[60] Yuval Tassa, Yotam Doron, Alistair Muldal, Tom Erez, Yazhe Li, Diego de Las Casas, David Budden, Abbas Abdolmaleki, Josh Merel, Andrew Lefrancq, et al. Deepmind control suite. arXiv preprint arXiv: /80100069,, 2018.
[61] Naftali Tishby, Fernando C Pereira, and William Bialek. The information bottleneck method. arXiv preprint physics/00040557 2000.
[62] Emanuel Todorov, Tom Erez, and Yuval Tassa. Mujoco: A physics engine for model-based control. In 2012 IEEE/RSJ international conference on intelligent robots and systems, pages 5026-5033. IEEE, 2012.
[63] Shenghua Wan, Yucen Wang, Minghao Shao, Ruying Chen, and De-Chuan Zhan. Semail: climinating distractors in visual imitation via separated models. In International Conference on Machine Learning, pages 35426-35441 PMLR, 2023.
[64] Tongzhou Wang, Simon S Du, Antonio Torralba, Phillip Isola, Amy Zhang, and Yuandong Tian. Denoised mdps: Learning world models better than the world itself. arXiv preprint ar0iv:2206. /5477, 2022.
[65] Manuel Watter, Jost Springenberg. Joschka Boedecker, and Martin Riedmiller. Embed to control: A locally linear latent dynamics model for control from raw images. Advances in neural information processing systems, 28, 2015.
[66] Daniel Willemsen, Mario Coppola, and Guido CHE de Croon. Mambpo: Sample-efficient multi-robot reinforcement learning using learned world models. In 2021 IEEE/RSJ International Conference on Intelligent Robots and Systems (IROS), pages 5635-5640. IEEE, 2021.
[67] Philipp Wu, Alejandro Escontrela. Danijar Hafner, Picter Abbeel, and Ken Goldberg. Day- dreamer: World models for physical robot learning. In Conference Robot Learning, pages 2226-2240. PMLR, 2023.
[68] Yuqin WU, Conggi SHEN, Shuhan CHEN, Chunming WU, Shunbin LI, and Ruan Wei. Intelli- gent orchestrating of iot microservices based on reinforcement learning. Chinese Journal of Electronics, 31(5):930-937, 2022.
14

---

[69] Xuantang Xiong, Linghui Meng, Jingging Ruan, Shuang Xu, and Bo Xu. Unec: Unsupervised exploring controllable space. In ICASSP 2024-2024 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP), pages 191-195. IEEE, 2024.
[70] Denis Yarats, Rob Fergus, Alessandro Lazaric, and Lerrel Pinto. Mastering visual continuous control: Improved data-augmented reinforcement learning. arXiv preprint arXiv:2107.09643 2021.
[71] Denis Yarats, Rob Fergus, Alessandro Lazaric, and Lerrel Pinto. Reinforcement learning with prototypical representations. In International Conference on Machine Learning, pages 11920-11931. PMLR, 2021.
[72] Denis Yarats, Ilya Kostrikov, and Rob Fergus. Image augmentation is all you need: Regu- larizing deep reinforcement learning from pixels. In International conference on learning representations. 2020.
[73] Denis Yarats, Ilya Kostrikov, and Rob Fergus. Image augmentation is all you need: Regu- larizing deep reinforcement learning from pixels. In International conference learning representations, 2021.
[74] Zhecheng Yuan, Guozheng Ma, Yao Mu, Bo Xia, Bo Yuan, Xueqian Wang, Ping Luo, and Huazhe Xu. Don't touch what matters: Task-aware lipschitz data augmentation for visual reinforcement learning. arXiv preprint 22Kiv:220220998222 2022.
[75] Amy Zhang, Rowan McAllister, Roberto Calandra, Yarin Gal, and Sergey Levine. Learning invariant representations for reinforcement learning without reconstruction. arXiv preprint 10742,, 2020.
15

---

A Environment
Walker
Quadruped
Jaco
Figure 7: We present visualization results for the three domains employed, following the introduction of driving car distractors. Notably, we retain the floor for Walker while removing the floor for Quadruped. This is due to the fact that Quadruped's floor occupies the entire observation space, precluding the addition of moving distractions without floor removal.
driving car
random video
dme
driving
dme
random_ video
Figure 8: To facilitate a clear understanding of the differences between the Driving-car and Random- video datasets, present visualizations of both datasets. Additionally, showcase the DMC observations of various distractor types. A detailed description of the datasets employed can be found in Section 5
To elucidate the direct impact of video distractors, we present illustrative visualizations in Figure 7 Various tasks present unique challenges. For instance, the multi-jointed Jaco arm with sparse rewards (only awarded for catching the red ball) poses difficulties in accurately reconstructing its dynamics model, hindering the learning process. Additionally, the absence physical floor and the introduction of complex distractor videos in quadruped tasks make it challenging for the agent to determine its relative position without relying on the floor. Additionally, Figure highllgght the distinct characteristics of the Driving Car and Random Video datasets. Videos in Random-vides feature RGB backgrounds that contrast sharply with the grayscale backgrounds in the Driving-car dataset.
B Proof
We base our approach on the following theoretical results from the paper.
16

---

B.1 Optimization objective of dynamics model
We define information bottleneck objective [6] to optimize the latent dynamics models same as many previous works [20 21 49 63]. The objective can be written as:
max 2=2::3 (Owr,, F1. (1rr) (a1:r) -8 [Z1 hn:r
(12)
(b)
Note that is a scalar and denotes indices of the dataset such that p(oy (i:) 8(0, o,) 22. We
constrain the capacity of information contained in latent state 21:7, while require them to predict
observations and rewards as accurate as possible. We can lower bound term (a) by the non-negativity of the KL divergence:
I(ai: (Our, P1. ((Ir)) Bptor moneemen)))))))))))))) In p Par Ra,r, 01.)) In p(1rr, Pi. Pa.l/a1:2) (1:r)
model-inelevant
(13)
t pplour mur (a [mp(orr, Pur rar (1r))
(14)
Beter In p (1:))
(15)
KL. Mirler, aur)l II gl
(16)
Σ In (orl2))))))))))))))))))))))))) + [no(orla)+ In(()))))
(17)
T
Σ
E a lenama 10ate lena) In q(oels s5 In (oneg(((()))) (st]
(18)
The second of Equation (3) can be dropped because the marginal data probability irrevelant to the dynamics model M. Equation (18) is obtained by the assumption of EX-BMDP that the latent state can be decoupled into endogenous part s; and exogenous part s;
For term (b), we can directly make use the derivation in SeMAIL 63 for its Equation (3), that
[[21r h (1r) $ B lont 1ne (o
(19)
(KL(po (sf 1st a o7 Il (sf |st ax-1)) + kI (s7 18 (8), l1go (st 88-11 at
1))
(20)
B.2 Proof of Proposition 4.1
Proposition 4.1 Restated. (Single step's mutual information) If only find the policy Texpl that max- imize every single step's 's mutual information ET= I(h st ax), then the whole trajectories' mutual information I(hī-ri El Teept) will be also maximized. Proof.
I(hiar' Elst Texp))
Σ
I(ht ( clh ht ht stt Teapl)
(21)
I(ht <<tt:<hh gl
ht h st ~ po(.1 ( ),
1T). Tespi)
(22)
z(nt:
El ;
Texp))
(23)
Σ
I(h aa)
(24)
0
17

---

Equation 22 is derived from the Chain Law of Mutual Information. Since we have already established the exploration policy and endogenous inference model pe. we can sample the endogenous state s+ for the next step. By iterating this process. we accumulate sufficient information to infer the endogenous state 8 as described Equation 23
C
Additional results
C.1 Visualization of observational trajectories
origin observations
O3
O
Og
06
0g
0g
03
0g
6g
0g
observation qu(oglst.sp
endo-decoder (g-(s()
Presentation reconstruction results from the observation model and endo-decoder. The first row shows the original observations, while ô1:4 represents the output of the observation model and of. represents the output of the endo-decoder. Each corresponding index pair indicates a one-to-one relationship between the original observation and the reconstructed observation.
Figure 9:
6+
6
mask
Figure 10: We present the following: the original observation o, images reconstructed with both the endogenous and exogenous branches ô, images reconstructed using only the endogenous branch öt, images reconstructed using only the exogenous branch ò and images with masking applied.
As illustrated in Figure 9 the observation model effectively reconstruets the original observations, including both background distractors and the agent. This indicates that s and capture sufficient information from observations. Furthermore, the endo-decoder outputs successfully reconstruct the endogenous information while eliminating exogenous distractors, demonstrating the ability of our separated model distinguish task-relevant information from task-irrelevant information, enabling the agent to make decisions without being influenced by extraneous factors.
18

---

In Figure 9 we have shown the reconstruction trajectories of SeeX after tinettunig in the Walker environment, but there was a lack of trajectory visualization of other environment during the explo- ration phase. Therefore, Figure [ displlay the reconstruction trajectories of SeeX during exploration in the Quadruped and Jaco environments. The trajectories are organized into five columns: from top to bottom, they represent the original observation the reconstruction ó obtained from the joint use of st and 5 with the help of the mask m,; the endogenous reconstruction 6+ obtained using only s+ representing the lask-relevant part; the exogenous reconstruction ô obtained using only representing the task-itrelevant moving distractors; and the mask m used reconstruct the entire image. By analyzing the reconstruction trajectories,
draw the following conclusions:
Although moving distractors were not completely removed, the reconstruction results show that s* contains most of the task-relevant information (including the agent's torso), significantly reducing the impact of background noise.
During the exploration phase, we use intrinsic rewards to encourage the agent to explore a more diverse range of states, aiming to build as comprehensive world model possible. This foundation supports imagination during policy training the fine-tuning phase. The reconstruction results indicate that the agent does not stick to fixed actions but attempts explore diverse states, demonstrating the effectiveness of the intrinsic rewards used.
The Jaco agent multi-joint robot with high flexibility, making it more complex than both the Walker and Quadruped. The reconstruction quality for Jaco is noticeably worse than for Quadruped, and the removal of moving distractors is less effective, which aligns with the experimental results in our paper. However, some interesting observations can be made: (1) The red ball, while not part of the agent's torso, is related the reward function. It appears ô+ but not in ô demonstrating SeeX's 's ability to automatically identify task-relevany information. (2) The base of Jaco, though fixed and not controlled by actions, part of the agent's torso but has little relation the reward. Therefore, SeeX classifies it as task- irreeeneeee,,,, corresponding to the reconstruction image ô=.
Average Performance on Different Tasks (Pre-training for 100k frames)
SeeX
Pxpnoee
APS
lormalized Return
0.8
APT
DIAYN
ICM
0.6
RND
0.4
0.0
Walker
Quadruped
Jaco
Fine-tuning performance comparison between SeeX and baselines, with 100k pre-training frames. Each bar represents the average normalized return across four walker tasks within a specific domain. We reported the mean of the top 30 highest returns.
Figure 11:
C.2 Complete pre-training performances
Owing to space constraints, the complete pre-training results are provided in the appendix. Figure 11 Figure [2] Figure [3] Figure presents the fine-tuned performance of pre-trained models trained with varying pre-training frame durations.
D Implementation Details
D.1 Networks and Hyper-parameter,
Building upon the recurrent state space model (RSSM) architecture 20 54 dynamics, posterior encoder, and policy components from the implementation in [54]. The hyper-
we retain the forward
19

---

Average Performance on Different Tasks (Pre-training for 500k frames)
1.0
Seek
Plan2 Explore
APS
Normalizee Return
0.8
APT
DIAYN
KCM
RND
0.4
0.0
Welker
Quadruped
Jeco
Fine-tuning performance comparison between SeeX and baselines, with 500k pre-training frames. Each bar represents the average normalized return across four walker tasks within a specific domain. We reported the mean of the top 30 highest returns.
Figure 12:
Average Performance Different Tasks (Pre-training frames)
1.0
Plan2Explore
APS
Normalized Return
DIAYN
0.6
LBS
0.2
Walker
Quadruped
Figure 13: Fine-tuning performance comparison between SeeX and baselines, with Im pre- trainininn frames. Each bar represents the average normalized return across four walker tasks within a specific domain. We reported the mean of the top 30 highest returns.
parameters for SeeX are detailed in Table We introduce masked image encoder, a novel image decoder for reconstructing the original observation, and an additional RSSM to model the transition of exogenous states.
Elaborating further, the exo-encoder Eg- mirrors the structure of the endo-encoder Egt (i.e., the image encoder in (54) However, introduce modifications the conventional decoder construct our novel endo-decoder Dø+ and exo-decoder Dø-, expressed in the following:
fully connected layer with 1536 hidden dimensions.
Reshape tensor into shape of (-1, 1536, 1).
transposed convolution layers with output channels, stride 2, and ELU activation with a 1.0.
Decompose 6-channel input into two 3-channel structures. The first represents the recon- struction output, while the remaining three channels encode the corresponding mask.
output transformation layer of Gaussian distribution.
To effectively reconstruct the observed state, introduce novel network architecture for the observation model (0:155; ;,5). This network integrates the outputs of the two decoders described carlier to generate the final reconstruction. The detailed network architecture is follows:
20

---

Average Performance on Different Tasks (Pre-training for 2m frames)
1.0
Seek
Plen2 Explore
APS
Normalized Return
0.8
APT
DIAYN
KCM
RND
0.4
0.0
Welker
Quadruped
Jeco
Figure 14: Fine-tuning performance comparison between SeeX and baselines, with 2m pre- pre-training frames. Each bar represents the average normalized return across four walker tasks within a specific domain. We reported the mean of the top 30 highest returns.
Get two images and two masks from Da and Do input.
2D convolution layer with 3 output channels and kernel size 1. Input: two masks, Output: a final mask.
Weight two images with the final mask as weights.
Output the final image.
Hyperparameter Value
& of Exo-Rec
3 for Walker, for Quadruped and Jaco
KL weight
1 for st and s
Deterministic size
Stochastic size
32
Discrete
32
Embedding size
200
Sequence length T 50
Batch size
50
Imagine horizon //
Precision
16
Model optimizer
Adam: (Ir: 3e-4, eps: le-5, clip: 100, wd: le-6)
Actor optimizer
Adam: (Ir: 8c-5, eps: 1e-5, clip: 100, wd: 1c-6)
Critic optimizer
Same as actor's
Action repcat
2
Table 2: Hyperparameters of SeeX for all experiments.
D.2 Detailed Description about Baselines
Curiosity is quantified as discrepancy between predicted consequences the agent's actions and the actual outcomes, represented in learned visual feature space. This formulation enables efficient exploration, outperforming baseline methods VizDoom and Super Mario Bros. Notably, ICM [46] exhibits superior exploration capabilities even in the absence of explicit external rewards. The study underscores the significance evaluating generalization to new scenarios and offers valuable insights into the effectiveness of reinforcement learning algorithms in adapting to novel environments.
ICM.
During exploration, Plan2Explore 54 effectively formulates plans to seek out novel experiences, enabling it to rapidly adapt to downstream tasks in zero-shot or few-shot setting.
Plan2Explore
21

---

Surpassing existing self-supprvised exploration techniques, Plan2Explore achieves competitive zero- shot task performance and eventually outperforms a supervised agent in few-shot scenarios. By leveraging unsupervised exploration to learn a world model, Plan2Explore demonstrates remarkable scalability and data efficiency, paving the way for significant advancements in building real-world reinforcement learning systems.
RND. By leveraging neural network to predict features of observations extracted from a randomly initialized network, RND [TO] effectively enhances exploration in Atari games characterized by sparse rewards. Notably, the method surpasses sfate-of-the-ar performance like Montezuma's 's Revenge, achieving remarkable results without relying on human demonstrations access to the game's 's underlying state. The paper underscores the flexibility of combining intrinsic and extrinsic rewards, demonstrating that RND facilitates directed exploration at a local level but faces challenges in global exploration over extended time horizons. Overall, the RND method represents significant progress in tackling hard exploration tasks, highlighting the potential of simple yet effective techniques at scale.
challenging games
LBS. Latent Bayesian Surprise (LBS) leverages Bayesian surprise in latent space to generate intrinsic rewards for exploration in Reinforcement Learning (RL). LBS surpasses existing methods on a diverse range of tasks, demonstrating efficient exploration in both continuous-control and discrete- action environments. It exhibits remarkable resilience to stochasticity in the environment dynamics, facilitating in-depth exploration and achieving superior performance high-dimensional settings such as video games. Notably, LBS significantly reduces computational costs while enhancing exploration capabilities, making a promising method for enhancing RL algorithms.
APT. APT 39 actively acquires behaviors and representations by exploring novel states in reward- free environments. This approach maximizes non-parametric entropy in abstract representation space, circumventing the need for complex density modeling and enabling effective scaling to high- dimensional observation environments. Leveraging intrinsic motivation and particle-based entropy maximization, APT achieves human-level performance on Atari games and surpasses benchmark results on DMControl tasks. Future research directions include reducing sample complexity through model-based RL. integration and investigating methods to mitigate catastrophic forgetting during fine-tuning. Overall, the paper showcases APT's efficacy in enhancing performance on challenging RL tasks while requiring substantially fewer samples compared to fully supervised RL algorithms.
DIAYN. By maximizing an information-theoreti objective using maximum entropy policy, DIAYN 16) enables the unsupervised emergence of diverse skills, such as walking and jumping, in simulated robotic tasks. The method demonstrates remarkable effectiveness in exploring complex environments, often solving benchmark tasks without receiving explicit task rewards. Additionally, DIAYN offers strategies for rapid task adaptation, hierarchical reinforcement learning, and imitation learning. By maximizing the mutual information between states and skills, DIAYN enhances the empowerment of hierarchical agent and ensures skill diversity through maximum entropy policies. Overall, the paper's 's contributions include proposing a novel method for unsupervised skill learning, showcasing the emergence of diverse skills in various tasks, demonstrating adaptability new tasks, and providing a foundation for hierarchical reinforcement learning and imitation tasks. DIAYN presents promising approach for enhancing exploration and data efficiency in reinforcement learning settings.
APS. By seamlessly integrating variational successor features with nonparametric entropy maxi- mization, APS (38) effectively optimizes the mutual information between tasks and policy-induced states. This method outperforms existing benchmarks on the Atari 100k data-cfficiency benchmark by employing entropy maximization to navigate the environment and leveraging data for behavioral learning. To stabilize training and improve convergence, the method incorporates averaging over k nearest neighbors. APS addresses the shortcomings of previous methods by maximizing the entropy of policy-induced states in a lover-dimensional abstract representation space. Empirical results on the Atari benchmark demonstrate state-of-the-arr performance, showcasing substantial progress over earlier work. The paper highlights the benefits of utilizing state entropy maximization data for las--conditioned skill discovery and outlines future research directions for further optimization and the integration of diverse approaches.
22

---

D.3 Experiments Compute Resources
We implement SeeX with Pytorch and run all the experiments on NVIDIA RTX 3090 for about 7500 GPU hours. The pre-training phase takes 10G of memory and 80 GPU hours for 2M frames, while the fine-tuning phase takes 10G of memory and 6 GPU hours for 100k frames. The codes of SeeX can be found in the supplementary materials.
NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: Our theoretical analysis and experimental results in Section andd Appendix C provide compelling evidence of the contributions and scope of our work.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification: A dedicated subsection Section within the "Contributions" section is devoted to discussing the limitations of our work.
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]
Justification: Detailed proof' of all assumptions and propositions presented in this paper be found the Appendix.
4. Experimental Result Reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that affects the main claims and/or conclusions the paper (regardless whether the code and data are provided or not)?
Answer: [Yes]
Justification: To facilitate reproducibility and further exploration, we provide comprehensive supplementary materials in the appendix, including a detailed description of the network framework Appendix D.I. and a comprehensive guide to hyperparameter settings Table 2
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
Justification: We provide the raw code and utilized datasets in supplementary materials.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary understand the results?
Answer: [Yes]
Justification: We show detailed experimental setting in Section 5 and corresponding visualization of environments and distractor videos in Appendix A Moreover, we provide descriptions of each baseline method in Appendix D.2
7. Experiment Statistical Significance
23

---

Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [Yes]
Justification: In the main performance results Figure 3] we present the standard deviation, while confidence intervals are shown for the Jaco task results Figure
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Question:
Answer: [Yes]
Justification: We show the detailed information of compute resources and relevent memory in Appendix D.3
9. Code Of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurlPS Code of Ethics
Answer: [Yes]
Justification: We have reviewed the NeurIPS Code of Ethics and confirm that our submission complies with the guidelines outlined in the document. We have considered the potential harms and societal impacts of our research and have taken steps to mitigate these risks.
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [Yes]
Justification: Unsupervised visual denoising reinforcement learning holds promise for en- hancing image quality with potential benefits in medical imaging, security, and historical preservation. However, societal risks such as deepfakes, algorithmic bias, and privacy concerns must be carefully considered. Mitigation strategies include developing deepfake detection methods, implementing fairness checks, and exploring privacy-preserve tech- niques. Overall, responsible development and deployment are crucial to ensure that this technology benefits society while minimizing potential harm.
Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: Our work on unsupervised visual denoising reinforcement learning does not involve high-risk data models. The focus is on improving image quality without generating or manipulating sensitive information. Therefore, safeguards for responsible release are not applicable in this context.
12. Licenses for existing assets
Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
The assets used in this paper, including code, training environments, and distractor videos, are all obtained from publicly available sources, their respective licenses and terms of use have been strictly adhered to. The code for this paper is built upon the code from [50]: the training environment adapted from [60], with MuloCo [62 as the underlying simulator; the distractor videos obtained from 227
Justification:
13. New Assets
24

---

Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification: Our paper does not release new assets.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification: The paper does not involve crowdsoureing nor research with human subjects. 15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: The paper does not involve crowdsourcing nor research with human subjects.
25


Reviews:
[{"review_id": "68DxIEHyc5", "reviewer": "Reviewer_SZqk", "summary": "The authors propose a method for separating endogenous and exogenous latent states for unsupervised exploration under visual distractors. Motivated by a theoretical bound minimizing the regret under a latent world model, the algorithm learns both an endogenous and exogenous world model, as well as an exploratory policy trained to maximize an intrinsic reward.\n\nThey evaluate on continuous control tasks with various distractors, showing their method outperforms other baselines. They perform additional experiments on how pretraining affects downstream finetuning RL performance and how their choice of intrinsic reward is a good approximation for the total variation found in the theoretical bound. They also provide analysis on parameter choices for their algorithm.", "questions": "Can you see this method helping in something like a self driving or navigation task in the real world?", "limitations": "Yes, they discuss the important limitation that the endogenous and exogenous state are always fully separate and don't interact with one another.", "rating": 8, "confidence": 3, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "The algorithm is extremely well-motivated from theory and does a good job arguing that the assumptions made in theory match reality. The paper is clear and rich with useful insights and the experimental results are strong.", "weaknesses": "Since the method performed so well in these settings, it would have been nice to see how it performed on a task in the real world under realistic endo/exo genous conditions", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your detailed review and valuable feedback on our work. We are very pleased that you found our algorithm theoretically motivated and experimentally successful. In response to your questions and suggestions, we have made the following points:\n- **Exploration of Real-World Applications**\n  - Our work focuses on handling moving distractors, where we need to separate task-relevant states from task-irrelevant states. In this setup, there are many common downstream tasks. For example, consider a book-finding robot in a library that needs to scan the shelves to find a specific book. The robot will encounter task-relevant objects (such as books) and task-irrelevant objects (such as posters on the wall or passing people), so the robot needs to distinguish which information is beneficial for the task.\n  - We agree with your view that real-world applications are crucial for validating the practicality of the algorithm. However, the self-driving and navigation challenges you mentioned are difficult to analyze and computationally expensive in practice. We would be willing to address these in future work. And this will involve designing experiments closer to real-world conditions to verify the algorithm's effectiveness in handling complex real-world scenarios.\n- **Discussion on Method Limitations**\n  - The assumption of completely separating endogenous and exogenous states that you mentioned is indeed a limitation of our method. We have discussed this in the paper, and our work has preliminarily validated the effectiveness of the bi-level optimization framework and the separated world model under the setting where endogenous and exogenous information is fully separable. In future work, we will explore how to introduce interactions between states to better understand and improve the applicability of the algorithm. \n\nOnce again, thank you for your support and constructive feedback on our work. We believe these improvements will further enhance the quality and impact of the paper.", "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": "q91QWOrQ4X", "reviewer": "Reviewer_wcfo", "summary": "This paper considers the problem of unsupervised reinforcement learning (task-agnostic pretraining) from image observations in environments with visual distractors. The key technical contribution of this paper is a practical algorithm, SeeX, based on the world model framework common in MBRL literature. SeeX explicitly separates endogenous (agent has control over) and exogenous (agent has no control over) forward dynamics in the learned world model, such that downstream policy optimization can occur only with the endogenous branch of the world model. The intuition behind the approach is that, assuming that the agent has no control over visual distractors (i.e., actions do not affect the distractors), we can expect such information to be encoded only in the exogenous branch of the model, thus not affecting policy optimization. The proposed method is validated on URLB, a common benchmark for unsupervised RL from images, and compared against prior work in this area.\n\n**Post-rebuttal comment:** I believe that the authors have addressed my concerns and I am revising my score from 5 -> 6 and recommend acceptance.", "questions": "I would like the authors to address my comments in the \"weaknesses\" section above. While I do have concerns regarding baselines, I believe that strong verbal arguments and/or changes to the writing is sufficient to address my concerns during the rebuttal. However, I will expect the authors to follow through with any promises of additional experimental results or changes to the paper for a future camera-ready version.", "limitations": "The authors discuss limitations in L331-334. However, I would like the authors to elaborate on the limitations of their approach.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "- The problem is interesting and relevant to the NeurIPS community, as well as RL researchers and practitioners more broadly. I believe that the approach is fairly straightforward yet original.\n- Paper is generally well written and easy to follow. There is sufficient background for an unfamiliar reader to appreciate the technical contributions.\n- Theoretical analysis helps build an intuition for the proposed method. Experiments are well thought out and the proposed method is compared to a variety of strong baselines from the URLB. Results seem to indicate that the proposed method indeed benefits from its separate branches.", "weaknesses": "I do not have any major concerns with the paper in its current form. However, there's a couple of things that I had expected to see in the paper:\n- A very common and effective strategy for visual generalization in RL has been use of data augmentation. There are numerous papers on the topic, many of which leverage simple and highly general augmentations such as random shift/crop, color jitter, mixup/overlay. Based on the strong similarity in problem setting, I would expect data augmentation to work very well on the considered benchmark. Intriguingly, data augmentation seems to have the same limitations as the proposed method, in the sense that they cannot capture distractors that the agent can interact with (i.e., has some control over). I would like the authors to comment on the similarity between the two research directions, and potentially include some discussion on this in the paper. I would strongly suggest including 1-2 data augmentation baselines for the camera-ready version of the paper as well; I fully acknowledge that this may not be a reasonable ask for the rebuttal itself.\n- I believe that there is another limitation of the proposed method which I didn't see any mention of in the paper. As far as I understand, moving objects that are task-relevant but external to the agent may not be captured by the endogenous branch. For example, a multi-agent setting in which another agent may move independently, or perhaps an environment in which a task-relevant object is moving due to e.g. physics without the agent actively interacting with the object. Is my understanding correct? This seems somewhat distinct from the example given in L333.\n- I understand why the authors would choose to only optimize the policy using the endogenous branch. However, I believe it would be quite valuable to include an ablation in which the policy is optimized with access to both branches; it is not obvious to me whether this would have any notable impact on downstream task performance.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the detailed review and valuable suggestions regarding our paper. We are pleased to hear that you find our work to have clear technical contributions and to perform well in experiments. In response to the questions and suggestions you have raised, we provide the following answers:\n- **Comparison between our research direction and data augmentation (DA)**\n  - Our designed bi-level separation framework extracts task-relevant information, while the DA method extracts task-relevant representations through data augmentation. As you mentioned, there are indeed similarities between the two approaches. However, since there are few methods that directly use DA for exploration, we incorporated the DA method into both Plan2Explore and SeeX to investigate whether it enhances performance. This is to verify whether the DA method provides benefits in the moving distractor setting.\n  - We used the classic data augmentation technique of random shift (pad by 4 pixels) proposed in DrQ[1], applying it 4 times to the same image. To thoroughly evaluate the effectiveness of DA, we tested its impact during both the pretraining (p-DA) and finetuning (f-DA) stages on the walker task. The performance results are shown in **Pdf-Table-A**.\n  - By analyzing the data in the table, we can draw the following conclusions:\n    - Applying data augmentation (f-DA) on specific tasks is effective, as it can mitigate the impact of distractors and improve performance, which supports the reviewer's insight.\n    - However, for tasks with moving distractors, the separation design in SeeX performs better than data augmentation.\n    - Using data augmentation during the pretraining stage leads to a performance drop, possibly because it interferes with the learning of the world model. The specific reasons will be explored in future work.\n- **Further clarification on the content from L331-334**\n  - Based on whether they affect the agent's reward function and action, distractors can be categorized into four types: task-irrelevant + action-independent (the setting addressed in our work); task-relevant + action-independent (which I understand to be the multi-agent setting you mentioned); task-irrelevant + action-dependent; and task-relevant + action-dependent. The latter two are the settings I referred to on L331-334, where distractors can interact with agents and even influence their reward function.\n  - We believe the reviewer's understanding is correct and we appreciate you pointing this out. We will include the settings you mentioned in the camera-ready version.\n- **Ablation in which the policy is optimized with access to both branches**\n  - Firstly, by formalizing the URL problem as a minmax problem, we further implement it as a separated world model consisting of two branches: endogenous and exogenous. The exogenous branch extracts task-irrelevant information, which is unrelated to the task. Therefore, it is reasonable to use only the endogenous branch to optimize the policy.\n  - Secondly, I understand that you are referring to the ablation experiments for using both branches as inputs to the policy. To address your concerns comprehensively, we consider the following two scenarios: (1) using both exogenous and endogenous states simultaneously for policy optimization, and (2) merging the endogenous and exogenous branches into one, using a single encoder to extract the state $s$ for policy optimization.\n    - (1) We added the following comparative experiments: both $s^+$ and $s^-$ are simultaneously input to both the actor $\\pi(a|s^+,s^-)$ and the critic $v(s^+,s^-)$, with the action still applied to the endogenous branch. We compared the performance of this setup (SeeX-both-branch) with SeeX in the walker environment, and the specific performance curves are shown in **Pdf-Figure-A**. In all tasks, SeeX outperformed SeeX-both-branch, especially in the last three challenging tasks. This indicates that using only endogenous information is advantageous for tasks with moving distractors.\n    - (2) Using a single encoder to encode observations for training the policy $\\pi(a|z)$, which is the approach of Plan2Explore[2]. The superiority of SeeX over Plan2Explore is demonstrated in Figures 3 and 4 of our paper, showcasing the effectiveness of the separated world model.\n  - In summary, we believe that using only the endogenous branch for policy optimization is reasonable, and our experiments and the content in the paper have demonstrated this. Although using information from two branches under our assumption doesn't offer advantages, as noted by the reviewer, there are multi-agent settings where noise, despite action independence, can affect the agents' reward functions. In such cases, incorporating part of $s^-$ into policy optimization might be effective. This is an interesting research direction that we are glad to explore in future work.\n\nThank you again for your support and constructive feedback on our work. We believe these improvements will further enhance the quality and impact of the paper. We will ensure that the additional experimental results and any necessary changes are included in the camera-ready version of the paper.\n\n[1] Image Augmentation Is All You Need: Regularizing Deep Reinforcement Learning from Pixels\n\n[2] Planning to Explore via Self-Supervised World Models", "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": "hz0ISvukwq", "reviewer": "Reviewer_HYxU", "summary": "This paper studies the problem of intrinsic-driving exploration in visually distracted environments, in the context of unsupervised reinforcement learning (URL). To address the issue that intrinsic rewards might be biased by distractors, the authors propose a method (called SeeX) that separates exogenous and endogenous information by training separate world models. In this way, the pre-trained world will be able to focus mainly on the task-relevant part of the observation, and the intrinsic reward based on dynamic uncertainty will be more accurate. The proposed method is evaluated on 3 image-based continuous control domains with video distractors. The results show that SeeX outperforms other URL methods on some environments.", "questions": "Please see weaknesses above.", "limitations": "The limitations have been briefly discussed in Section 7.", "rating": 4, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "* Separating exogenous and endogenous information to produce a better intrinsic reward is an interesting idea.\n* The writing is mostly clear. The diagrams look nice.", "weaknesses": "(Not all points are weaknesses. Comments/questions/weaknesses are all gathered in this section for convenience.)\n\n* Line 157: what is BA?\n* Line 164: should the second $=$ be $+$?\n* Figure 4b (walker stand): this figure is a bit misleading. It seems to suggest that SeeX is much better than other methods for pretrain budgets between 100k ~ 500k. However, this is not necessarily true, as there are no data points at 200k/300k/400k for baselines. The claim that other methods require **at least** 500k frames to achieve the same level of performance is ungrounded.\n* Apart from Figure 9, could the authors provide more visualizations for the decoded observations? For example, $\\hat{o}_t^-$, $m_t$, and trajectories with different embodiments.\n* Have the authors conducted ablation experiments (K/Se/Exo-Rec) on other environments than Quadruped? Considering the quadruped results are quite noisy (Figure 3), it would be better to also include ablation results on walker envs.\n* Ablation: how does SeeX perform when $\\alpha=0$?\n\nOther issues:\n* The legend of Figure 5 is a bit confusing. It took me some time to understand that the bottom legend is for the bottom-right figure. Please consider rearranging its position.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your detailed review. We have noted the typographical and formatting issues mentioned by the reviewer. BA on Line 157 refers to the Barber-Agakov bound[1], named using the authors' initials, which is a commonly used bound for mutual information. The second \"=\" on Line 164 should indeed be a \"+\". We believe these adjustments do not affect the main content and conclusions of the paper. Below are our responses to the issues you raised.\n- **Regarding conclusions of Figure 4**\n  - Our original intention was to show SeeX's performance stability between 100k and 500k steps; we sampled and tested the performance of other methods between 100k~500k, but none exceeded 90% expert performance, so we did not report them in the paper. We show the performance of all steps for the other methods in **Pdf-Table-B** and our refined conclusion is that SeeX achieves 90% expert performance with fewer pre-training steps than other methods, demonstrating the efficiency of our bi-level framework and the separated world model during exploration.\n- **Other trajectory plots except for Figure 9**\n  - In Figure 9, we have shown the reconstruction trajectories of SeeX after fine-tuning in the Walker environment, but there was a lack of trajectory visualization of other environment during the exploration phase. Therefore, **Pdf-Figure-C** displays the reconstruction trajectories of SeeX during exploration in the Quadruped and Jaco environments. The trajectories are organized into five columns: from top to bottom, they represent the original observation $o$; the reconstruction $\\hat{o}$ obtained from the joint use of $s^+$ and $s^-$ with the help of the mask $m$; the endogenous reconstruction $\\hat{o}^+$ obtained using only $s^+$, representing the task-relevant part; the exogenous reconstruction $\\hat{o}^-$ obtained using only $s^-$, representing the task-irrelevant moving distractors; and the mask $m$ used to reconstruct the entire image. By analyzing the reconstruction trajectories, we can draw the following conclusions:\n    - Although moving distractors were not completely removed, the reconstruction results show that $s^+$ contains most of the task-relevant information (including the agent's torso), significantly reducing the impact of background noise.\n    - During the exploration phase, we use intrinsic rewards to encourage the agent to explore a more diverse range of states, aiming to build as comprehensive a world model as possible. This foundation supports imagination during policy training in the fine-tuning phase. The reconstruction results indicate that the agent does not stick to fixed actions but attempts to explore diverse states, demonstrating the effectiveness of the intrinsic rewards used.\n    - The Jaco agent is a multi-joint robot with high flexibility, making it more complex than both the Walker and Quadruped. The reconstruction quality for Jaco is noticeably worse than for Quadruped, and the removal of moving distractors is less effective, which aligns with the experimental results in our paper. However, some interesting observations can be made: (1) The red ball, while not part of the agent's torso, is related to the reward function. It appears in $\\hat{o}^+$ but not in $\\hat{o}^-$, demonstrating SeeX's ability to automatically identify task-relevant information. (2) The base of Jaco, though fixed and not controlled by actions, is part of the agent's torso but has little relation to the reward. Therefore, SeeX classifies it as task-irrelevant, corresponding to the reconstruction image $\\hat{o}^-$.\n- **Regarding the ablation studies in Figure 5**\n  - When formatting Figure 5, we only considered aesthetic factors and overlooked the potential for misleading the reviewers. Therefore, we have re-drawn the conclusions of the ablation study and improved the readability of the charts, and added the ablation experiments on the walker (**Pdf-Table-C**) as well as the performance with $\\alpha=0$ (**Pdf-Figure-B**). By analyzing the charts, we can draw the following conclusions:\n    - The choice of weight for the Exo-Rec term significantly impacts SeeX's performance. For Walker, selecting $\\alpha=3$ is appropriate, while for Quadruped and Jaco, $\\alpha=1$ is more suitable. Removing the Exo-Rec term (equivalent to setting $\\alpha=0$) leads to varying degrees of performance degradation across different environments, indicating that the Exo-Rec design contributes positively to performance.\n    - Observing **Pdf-Table-C**, SeeX's sensitivity to different hyperparameters shows similar conclusions for Walker and Quadruped: (1) Removing the separated world model significantly impacts performance; (2) Removing the Exo-Rec term aligns with the conclusions from **Pdf-Figure-B**; (3) SeeX's performance is generally positively correlated with the number of predictive heads $K$, but considering the computational cost associated with increasing $K$, $K=5$ is a suitable choice.\n\nThank you again for your support and constructive feedback on our work. We believe these improvements will further enhance the quality and impact of the paper.\n\n[1] The IM algorithm: a variational approach to information maximization", "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": "tExXXxObtx", "reviewer": "Reviewer_NjAE", "summary": "To address the challenge of visual distractions in unsupervised reinforcement learning, this paper proposes a bi-level optimization framework called SeeX, which utilizes a separate world model to mitigate the disturbance caused by visual distractions. The authors evaluate the proposed method in multiple tasks.", "questions": "- Line 193: \"It is also an approximation to the world model TV error in Problem 3.3\". Can the author provide more explanation on why Eq (11) is consistent with Eq (7)?\n- The authors wish to have a separate world model. However, how do you ensure that $s^-$ does not contain information about the reward? Since you have $\\hat o_t^{exo}$?\n- In the experiment, what RL algorithm are the baselines implemented with?\n- In Eq (6), the REGRET is bounded by the estimation error of $\\hat{\\mathcal{T}}^+$ and $\\hat{\\mathcal{T}}^-$. However, if $s^-$ is irrelevant to reward signals and endogenous state $s^+$, and $\\pi$ only depends on $s^+$, why would the estimation error of $\\hat{\\mathcal{T}}^-$ affect the REGRET?", "limitations": "N/A", "rating": 4, "confidence": 4, "soundness": 2, "presentation": 1, "contribution": 2, "strengths": "- While unsupervised RL relies on intrinsic reward to promote exploration, the existence of visual distractions might contaminate the evaluation of intrinsic reward. SeeX proposes to use a separate world model to mitigate such disturbance.", "weaknesses": "- The separation model assumes decoupled endogenous and extraneous transition, which could be a strong assumption. For example, in robot manipulation, obstacles might be irrelevant to reward signals but still affect the movement of the robot.\n- The separation model uses a mask model to combine endogenous and extraneous reconstruction. However, this might not work well when the distractions and embodiments are inseparable, e.g. Color Change or Camera Jittering [1].\n- The proposed EX-BMDP seems to be similar to Task-Informed MDP in [2], while the proposed intrinsic reward is similar to the latent disagreement in [3], I have some reservations about the novelty of combining these two methods as they are generally orthogonal.\n\n[1] Denoised MDPs: Learning World Models Better Than the World Itself\n\n[2] Learning Task Informed Abstractions\n\n[3] Planning to Explore via Self-Supervised World Models\n\nTypos:\n- Line 51we use the task-irrelevant (-> task-relevant) information to predict the reward and the action.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your review and feedback. The term \"task-irrelevant\" on L51 should be \"task-relevant\". Here are our responses to your questions:\n- **Strong assumption**\n  - Firstly, many Visual RL works use separation assumptions, such as Denoised MDP[3] and TIA[2], which have natural assumptions and practical significance. SeeX is inspired by this line of research. We believe no single assumption is universally applicable (e.g., CNNs for images and RNNs for text). So whether an assumption is \"strong\" is subjective and context-dependent.\n  - Secondly, we focus on URL problems with task-irrelevant and action-independent noise and we model URL as minmax problems through theoretical analysis. A **practical scenario** fitting our assumption is a book-finding robot in a library. The robot must differentiate between task-relevant items (e.g., books) and task-irrelevant items (e.g., posters or people) to find the specific book. Thus, our work is grounded in both theory and practical application.\n  - Finally, the robot manipulation you mentioned is discussed in L331-334.\n  \n- **Inseparable noises**\n  - We emphasized in the paper (caption of Figure 1 and the conclusion) that we only focus on handling moving distractors. Moving distractors and inseparable noises represent **distinct research directions**. According to the No Free Lunch theorem, no single method can be effective for all tasks. \n  \n- **Discussion on whether EX-BMDP and latent disagreement are orthogonal**\n  - **Separated world model** \n    - Firstly, we **adopted**(L78-79) rather than **proposed** EX-BMDP[1]. The similarity you see between EX-BMDP and TIA is due to their shared use of the separation concept, **but separation is not a specific method but rather a general idea**[1,2,3,4].\n    - Additionally, EX-BMDP differs significantly from TiMDP[2]. TiMDP applies action to the exogenous branch, while EX-BMDP does not. Did the reviewer notice this? The agent can only influence the endogenous branch under our assumption, which is why EX-BMDP is suitable.\n  - **Intrinsic reward**\n    - As we all know, during pretraining phase of unsupervised RL, task-specific rewards are typically avoided in favor of intrinsic rewards, which come in various forms[5,6,7,8]. **Any intrinsic reward can support SeeX's exploration**, and we chose the classic disagreement[5].\n  - Our method formalizes URL as a minmax problem, where the outer layer minimizes REGRET’s upper bound using a separation assumption (Eq 2), leading us to adopt EX-BMDP. Inner layer maximize state uncertainty, so we approximate Eq 7 with Eq 11, making disagreement a reasonable choice. Reviewer considers the methods to be orthogonal, but we believe that the underlying ideas abstracted (e.g., separation and intrinsic reward) can inspire cross-field applications. Our step-by-step design and component choices are based on necessity, and our results validate the framework's effectiveness.\n  \n- **L193, the relationship between Eq 7 and Eq 11**\n  - In Eq 7, the goal of $\\pi_{expl}$ is to maximize the total variation (TV) distance between the learned model and the true latent dynamics in the task-relevant part. Practically, we estimate $D_{TV}$ using the disagreement among an ensemble of neural networks, inspired by [9].\n  \n- **How to ensure $s^-$ does not contain information about the reward?**\n  - Firstly, our Exo-Rec design follows [2,3], which states that task-relevant information is a small part of the observation. The Exo-Rec term aims to encode as much information from $o$ into $s^-$ (L164-166), not to remove reward-related information from $s^-$ as you mentioned.\n  - Secondly, We assume that only the task-relevant part of $o$ is affected by the action. Based on EX-BMDP, we design a separated world model where the action limits the encoding of task-relevant information into $s^+$.\n  - In summary, **we believe there is a misunderstanding of our method and the Exo-Rec design**. In visual settings, reward-relevant and reward-irrelevant information are inherently mixed and cannot be perfectly separated. Are you asking us to ensure a 100% separation between the two? This is not feasible. Our results(Figure 3, Table 1) show that our method attempts to separate $s^-$ and $s^+$ as much as possible, benefiting the task of handling moving distractors.\n  \n- **The implementation of baseline methods**\n  - We have mentioned in the paper (L47-49, L73-77, L201) that our method builds on URLB[10] and URLB-pixel[11], using the same implementation for the baseline. The specific implementation of the baseline can be found in the URLB-pixel GitHub repository.\n  \n- **The REGRET bound in Eq (6)**\n  - According to your statement, if we extract **precise** task-irrelevant information ($\\hat{T}^-$) from the observation, it wouldn’t help in reducing REGRET? This seems unreasonable because reward-relevant and reward-irrelevant information are complementary and together constitute the observation. Knowing one allows us to infer the other.\n  - Although $s^-$ is unrelated to REGRET, we use losses of $s^-$ and $s^+$ to bound REGRET (Eq 6). If you believe that $s^-$ is insignificant, you might try deriving a tighter REGRET bound based solely on $s^+$.\n\nWe hope this clarifies our work. Please feel free to ask any further questions.\n\n[1] Provably Filtering Exogenous Distractors using Multistep Inverse Dynamics\n\n[2] Learning Task Informed Abstractions\n\n[3] Denoised MDPs: Learning World Models Better Than the World Itself\n\n[4] SeMAIL: Eliminating Distractors in Visual Imitation via Separated Models\n\n[5] Planning to Explore via Self-Supervised World Models\n\n[6] Curiosity-driven exploration by self-supervised prediction\n\n[7] Exploration by random network distillation\n\n[8] Curiosity-driven exploration via latent bayesian surprise\n\n[9] Reward-free curricula for training robust world models\n\n[10] URLB: Unsupervised Reinforcement Learning Benchmark\n\n[11] Mastering the unsupervised reinforcement learning benchmark from pixels", "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:
We thank the reviewers for taking the time to give useful comments for our paper. We are glad that the Reviewers appreciate the correct motivation (Reviewer SZqk), original approach (Reviewer wcfo) and strong experimental results (Reviewer SZqk, wcfo). Reviewers pointed out the concerns and points for improvement, and we made responses for each of them. We presented more experimental results and detailed responses to questions about our assumption and approach. \
We kindly request the reviewers to inform us if there is anything else we can clarify. Thank you for taking the time to review our work.

Meta-review: This paper proposes a new unsupervised RL algorithm for environments with visual distractors. The key idea is to learn a separated world model to distinguish between task relevant information and task irrelevant distractors. There is agreement that the approach is sound, motivations are clear, has a good theoretical/experimental grounding. 

The key strengths that were highlighted were: the approach is novel and well motivated, theoretical foundation and analysis is strong, the benchmark tasks are impressive, and the presentation is clear. Although there were some weaknesses reported as well: justification and validation for why certain assumptions were made, relationship to prior work and scope of evaluation outside of standard benchmarks, to more real world applications.

Reviewers answered detailed responses to questions raised during the review process, including ablation studies and clarifications. I recommend accepting this paper for a poster presentation.
Decision: Accept (poster)