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_cc_2024_Conference
Published: 2024-05-14 09:34:40

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.

Main paper: 
 intrinsic reward: compute endogenous uncertainty under separation view: 
 Crafting an effective intrinsic reward function is paramount during the pretraining (PT) phase. A welldesigned intrinsic reward can steer the agent towards exploring states with the highest uncertainty, maximizing skill acquisition. A common practice [44, 45, 52] is to maximize the mutual information I(h + 1:T ; ξ|s + 0 , π expl ), where ξ represents the true unknown dynamics. This objective is motivated by the fact that mutual information quantifies how comprehensively the trajectory explores the environment. Proposition 4.1. (Single step's mutual information) If only we find the policy π expl that maximize every single step's mutual information T -1 t=0 I(h + t+1 ; ξ|s + t , a t ), then the whole trajectories's mutual information I(h + 1:T ; ξ|s + 0 , π expl ) will be also maximized. Proof in Appendix B.2. With Proposition 4.1, the optimization objective can be written as:EQUATIONTo approximate the unknown ξ, we design a set of predictive heads {q k ( ĥ+ t+1,k |ξ k , s + t , a t )} k=1:K which are implemented as conditional Gaussians N (µ(Ξ = ξ k , s + t , a t ), σ 2 ). Given fixed variance, the conditional entropy does not depend on state or action [52] . Suppose that p(Ξ) is an uniform distribution on {ξ k } k=1:K , then we have marginal entropy in Equation ( 11) lacks a closed-form expression amenable to optimization, we employ the empirical variance across ensemble means as a substitution. It is also an approximation to the world model TV error in Problem 3.3. We define the intrinsic reward as r intr t = Var({ ĥ+ t+1,k } k=1:K ), and employ the exploratory policy π expl that maximizes r intr t as an approximation of π * expl in Equation (10) .H(h + t+1 |s + t , a t ) = 1 K K k=1 H(h + t+1 |Ξ = ξ k , s + t , 
experiments: 
 All experiments are conducted with at least three seeds and evaluated for 10 episodes. We conduct experiments to answer the following questions: 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-left). These tasks vary in difficulty, providing a 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 in Appendix A. Agents receive only visual inputs (64, 64, 3) , with episodes lasting 1000 frames and an action repeat of R = 2. 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 [27] , we construct two subdatasets for our experiments: the Driving-car dataset, commonly used in previous works [18, 61] , and the Random-video dataset, composed of 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 Random-video 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-based: ICM [44] , RND [10] , LBS [42] and Plan2Explore [52] maximize prediction error to better understand the world; Data-based: APT [38] encourages exploration by maximizing entropy; Competence-based: DIAYN [16] and APS [37] maximize mutual information to achieve diverse discovery and generalization. The implementation detail and more concrete introduction of the above methods are shown in Appendix D.2. Given the small gap at 500k on Walker-Stand, a table highlights details, with red marking over 90% expert performance.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-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 lowest performance exceeds the highest curves of other methods, a trend seen across quadruped tasks. Plan2Explore shows relatively weaker results, especially in the walker-stand and walker-walk tasks, likely due to relying solely on predictive heads based on 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. ) values. This suggests that the total environmental uncertainty remains unchanged; rather, most uncertainty is concentrated in q θ -and p θ -due to SeeX's effective capture of task-irrelevant distractor transitions.Building upon the promising results from the previous section, where SeeX outperformed seven baseline methods and achieved expert-level performance from [33] (trained in a distractor-free environment, representing the upper bound for our setting) in some tasks, we delve into an intriguing us values of 𝛼 More specifically, can we achieve fine-tuned performance with as few pre-training steps as possible to achieve more than 90% of expert performance? We conduct experiments on walker-stand and walker-walk. For all methods, we assessed the fine-tuning performance using pre-training frames of (100k, 500k, 1m, 2m). To evaluate SeeX on a more granular level, we add choices of (20k, 50k, 200k, 300k, 400k). As shown in Figure 4-(b ), for the Walker-Walk task, SeeX reaches 90% expert performance (red dotted line) within 500k frames. In contrast, no baseline achieves this level, even with 2M pre-training frames. For the simpler Walker-Stand task, a table offers detailed analysis. SeeX requires only 100k frames of random exploration to reach 90% expert performance, while other baselines need significantly more pretraining.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 5-(a) , removing either the separate world model or 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.In SeeX, we use only the endogenous state s + to predict actions. To validate this design, we compare two policies: π(s + ) (SeeX) and π(s + , s -) (SeeX-both-branch). As shown in Figure 5-(b) , SeeX consistently outperforms SeeX-both-branch, particularly in the last three challenging tasks. This demonstrates that relying solely on endogenous information is beneficial in environments with moving distractors.Ablation of α. We conducted experiments to assess the impact of Exo-Rec's weight α. As shown in Figure 5-(c ), the optimal value of α varies across domains. For Jaco and quadruped tasks, a smaller α (e.g., 1) enhances performance, whereas walker tasks benefit from a larger weight (e.g., 2 or 3) to effectively extract exogenous information into s -.In this subsection, we 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-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 6-(b) , SeeX outperforms baselines on test distractor datasets and shows minimal performance drop under distribution shifts, highlighting its strong generalization.In this subsection, we evaluate the effectiveness of SeeX in addressing the "generalize to distractors from other distributions" 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-tuned agents on the driving-car dataset and evaluated their performance on both drivingcar 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 Table 2 , SeeX consistently outperforms the baselines on both test distractor datasets and exhibits the smallest performance drop under distribution shift, demonstrating its clear advantages in generalization ability.Driving Owing to space constraints, the complete pre-training results are provided in the appendix. Figure 8 presents the fine-tuned performance of pre-trained models trained with varying pre-training frame durations.Building upon the recurrent state space model (RSSM) architecture [19, 50] , we retain the forward dynamics, posterior encoder, and policy components from the implementation in [50] . The hyperparameters for SeeX are detailed in Table 3 . We introduce a 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 E θ -mirrors the structure of the endo-encoder E θ + (i.e., the image encoder in [50] ). However, we introduce modifications to the conventional decoder to construct our novel endo-decoder D θ + and exo-decoder D θ -, as expressed in the following:• 1 fully connected layer with 1536 hidden dimensions.• Reshape tensor into shape of (-1, 1536, 1, 1).• 4 transposed convolution layers with 6 output channels, stride 2, and ELU activation with α = 1.0. • Decompose 6-channel input into two 3-channel structures. The first represents the reconstruction output, while the remaining three channels encode the corresponding mask.• 1 output transformation layer of Gaussian distribution.19 Numerous studies [3, 70, 67, 41] explore data augmentation (DA) in 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 to 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 (3) pDA reduces performance, likely by disrupting world model learning, further investigation is planned. 
related work: 
 Model-based Control. Learning a dynamics model of the environment is a 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 of URLB-pixels, which combines all baseline methods into a unified Dreamer-like framework. Plan2Explore makes use of the dreamerv2 framework and takes the uncertainty of prediction of the next latent state as the intrinsic reward. In our work, we 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.Unsupervised RL. In recent years, there are many works that tried to promote the unsupervised representation learning manner in various fields including computer vision (CV) [12, 23, 25] and natural language processing (NLP) [9, 46, 14] . These works encourage the RL community to explore the more efficient way of learning [34, 32, 51, 56, 69, 11, 26] , however, 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 RL has enabled agent for diverse exploration and data [38, 53, 68] ; (ii) Knowledge-based: increase knowledge about the world with self-supervised prediction [44, 45] ; (iii) Competence-based: methods based on mutual information [16, 22, 37, 55] shows capability for diverse discovery and generalization. URLB-pixels offers a unified model-based framework to implement some methods referred to above, we make use of this benchmark to serve as the baseline. A 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 endogenous latent space to train policy. As a 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 [69, 17, 71, 7] ; (2) Learning task-relevant representations with bisimulation metrics [72, 39] ; (3) Design auxiliary tasks to extract endogenous information [4, 5, 15, 31] ; (4) Take actions or rewards as discriminating factors to separate exogenous and endogenous information [18, 62, 43, 40] . 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 in complex real-world environments. 
conclusion: 
 In this paper, we 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 a separated world model and task-relevant uncertainty maximization to mitigate the impact of distractors and enhance exploration efficiency. In this paper, we 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 a separated world model and task-relevant uncertainty maximization to mitigate the impact of distractors and enhance exploration efficiency. Furthermore, 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 a separated world representation to 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. Extensive 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-independent, task-irrelevant + action-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-independent case (e.g., multi-agent systems), we propose incorporating some s -into policy optimization as a potential solution. 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 of a 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 8 highlights the distinct characteristics of the Driving Car and Random Video datasets. Videos in Random-video feature RGB backgrounds that contrast sharply with the grayscale backgrounds in the Driving-car dataset. 
: 
 Unsupervised learning has a rich history in computer vision and natural language processing, as demonstrated by methods such as [12, 23, 25, 9, 46, 14] . It leverages unlabeled, task-agnostic data to train models that can be quickly adapted to downstream tasks, addressing sample inefficiency. 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.A key advantage of URL is that learned dynamics models can gather prior knowledge about the environment during exploration, minimizing the need for extensive interactions during adaptation to downstream tasks.A major challenge in unsupervised reinforcement learning arises from distractors, which significantly hinder learning in complex real-world environments. For example, a book-finding robot in a 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 [62] . As shown in Figure 1 -Left, a single world model encodes both task-relevant and task-irrelevant information in the latent space. Even under a reward function R. Let V M R (π) denotes the expected value of policy π under dynamics M with reward function R. The regret of policy π is defined as REGRET(π, M R ) = V M R (π * ) -V M R (π), π * = arg max π ′ V M R (π ′ ), which measures the performance gap between π and the optimal policy. The optimization goal is to find the world model that is robust to different possible reward functions, minimizing the regret of world model policy under real dynamics:EQUATIONIn 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 a novel, reward-free objective using Simulation Lemma [28, 50] . Our proposition is outlined below.Proposition 3.2. 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:EQUATIONAs outlined in Proposition 3.2, the optimal world model policy exhibits low regret if the latent state-action distribution of both π * R and π * R in M is accurately captured by T . However, during unsupervised pre-training, the reward function remains inaccessible, preventing us from directly obtaining these distributions. This is because the state-action distribution induced by π * R , and π * R is inherently dependent on the reward function. To address this challenge, we introduce an exploration policy π expl , specifically designed to maximize the expected error (in terms of total variation distance) of the latent dynamics model:EQUATIONThis enables us to derive an upper bound on the regret that is independent of the reward function: (4)REGRET( π * R , M R ) ≤ 4γ (1 -γ) 2 E z,In the context of EX-BMDPs, the transition dynamics can be decomposed into two components: T (z ′ |z, a) = T + (s + ′ |s + , a)T -(s -′ |s -), where z = (s + , s -), s + represents the endogenous state, and s -represents the exogenous state. To apply EX-BMDP to visually distracted environments, we further assume that the reward depends solely on s + , implying that decisions are made based only on endogenous state.To model this separation, we construct the world model as two distinct components: the endogenous part M + and the exogenous one M -, corresponding to the task-relevant and task-irrelevant processes, respectively. Their associated latent dynamics are denoted by T + and T -. Utilizing this decomposition, we can derive a lower bound for the policy optimization target:EQUATIONThe strategy's optimization objective is a form of model uncertainty. Estimating model uncertainty for the URL serves as a reward for the exploration policy, which is crucial for learning. However, under the EX-BMDP assumption, directly maximizing the reward of 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:EQUATIONWith these bounds, we can reformulate the problem as a bi-level optimization problem: Problem 3.3. (Bi-level Optimization of World Model Error) Consider a reward-free EX-BMDP with latent dynamics functions T + , T -and an exploration policy π, assuming the reward depends only on the endogenous state. Let T = T + × 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 Observation Model. Some model-based approaches [20, 47] use an auxiliary reconstruction loss to refine the observation encoder, ensuring z t captures sufficient information from o t by maximizing the mutual information I(o t ; z t ). The IM algorithm [6] provides the Barber-Agakov lower bound:EQUATIONI(o t ; z t ) ≥ E p(ot,zt) [ln q θ (o t |z t )]+H(p(o t )), where H(p(o t )) is the entropy of the observation distribution. In our separation setting, this bound becomesI(o t ; s + t , s - t ) ≥ E p(ot,s + t ,s - t ) [ln q θ (o t |s + t , s - t )].Inspired by [61, 18, 24] , We design an observation model q θ (o t |s + t , s - t ), an endo-decoder(ô + t , m - t ) ∼ D θ + (s + t ), an exo-decoder (ô - t , m - t ) ∼ D θ -(s - t ) and a mask-mixing model m t = M θ (m + t , m - t ). ô+ t and ôt are expected to reconstruct the task-relevant and task-irrelevant parts of the original observation, while m + t and m - t are corresponding masks. So q θ (ot |s + t , s - t ) can be implemented as ôt = m t ⊙ ô+ t + (1 -m t ) ⊙ ô- t .Considering insights from [18, 62] that task-relevant information occupies only a small portion of the observation, we design an exogenous reconstruction (Exo-Rec) model ôexo t ∼ q θ (o t |s - t ) to ensure that s -contains the vast majority of information.World Model Optimization. Similar to a variational autoencoder (VAE) [30, 49] , all model components are trained jointly by maximizing the evidence lower bound (ELBO). Our optimization objective extends that of Dreamer [20] :𝑜 ! 𝑜 !"# exogenous branch endogenous branch 𝑠 ! " 𝑠 !"# " 𝑠 ! $ 𝑠 !"# $ 𝑎 !"# 𝑠 ! " 𝑂 ! 𝑠 # " # 𝑣 ! r ! 𝑎 ! # 𝑣 ! r ! 𝑎 ! 𝜋 𝜋 (b) outer optimization of policy intr intr 𝑟 ! "#$% 𝑟 !&' "#$% 𝑠 % " 𝑎 % concat ℎ ! ! ℎ !"#,# " ! ℎ !"#,% " ! ℎ !"#,& " … 𝑟 " #$%& = 𝑉𝑎𝑟( ' ℎ "!',) ! )*':, ) 𝜉 ' 𝜉 - 𝜉 , 𝐸 ( ! 𝐸 ( ! 𝐸 ( " 𝐸 ( " % 𝑜 ! %&' % 𝑜 ! % 𝑜 !"# %&' % 𝑜 !"# 𝑎 ! … … … … … imaginary trajectories (a) inner optimization of world model ℎ !"# $ ℎ ! $ ℎ !"# " ℎ ! "L . = E p ( t (L t O + L t R + L t KL )).The detailed formulations of each term are presented as follows, and the derivations are shown in Appendix B.1.EQUATIONEQUATIONL t KL . = KL(p θ + (s + t |s + t-1 , at-1, h + t )||q θ + (s + t |s + t-1 , at-1)) + KL(p θ -(s - t |s - t-1 , h - t )||q θ -(s - t |s - t-1 , 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 π(a t |s + t ) and a reward model q θ (r t |s + t ). The reward model fits the true reward function in the FT phase. The actor selects action based on the current endogenous state s + t under the guidance of intrinsic reward in the PT phase or the reward model in the FT phase. 
 
 Appendix: Walker Quadruped JacoWe base our approach on the following theoretical results from the paper.We define an information bottleneck objective [59] to optimize the latent dynamics models same as many previous works [20, 21, 47, 61] . The objective can be written as: (12) Note that β is a scalar and i t denotes indices of the dataset such that p(omax I(z 1:T ; (o 1:T , r 1:T )|a 1:T ) (a) -β I(z 1:T , i 1:T |a 1:T ) (b)t |i t ) = δ(o t -o ′ t ) [2]. We constrain the capacity of information contained in latent state z 1:T , 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: -KL p(o1:T , r1:T |z1:T , a1:T )|| T t=1 q(ot|zt)q(rt|zt)= E q(z 1:T |o 1:T ,a 1:T ) T t=1 ln q(ot|zt) + ln q(rt|zt)EQUATIONThe second term of Equation ( 13) can be dropped because the marginal data probability is irrevelant to the dynamics model M. Equation ( 18) is obtained by the assumption of EX-BMDP that the latent state z t can be decoupled into endogenous part s + t and exogenous part s - t . For term (b), we can directly make use of the derivation in SeMAIL [61] for its Equation (3), that t=0 I(h + t+1 ; ξ|s + t , a t ), then the whole trajectories's mutual information I(h + 1:T ; ξ|s + 0 , π expl ) will be also maximized. Proof.EQUATIONEQUATIONEQUATIONEQUATIONEQUATIONEquation ( 21) is derived from the Chain Law of Mutual Information. Since we have already established the exploration policy and endogenous inference model p θ , we can sample the endogenous state s + for the next step. By iterating this process, we accumulate sufficient information to infer the endogenous state s + t-1 , as described in Equation ( 23).C. As illustrated in Figure 9 , the observation model effectively reconstructs the original observations, including both background distractors and the agent. This indicates that s + t and s - t capture sufficient information from the observations. Furthermore, the endo-decoder outputs successfully reconstruct the endogenous information while eliminating exogenous distractors, demonstrating the ability of our separated model to distinguish task-relevant information from task-irrelevant information, enabling the agent to make decisions without being influenced by extraneous factors.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, Figure 10 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 ô obtained from the joint use of s + and s -with the help of the mask m; the endogenous reconstruction ô+ obtained using only s + , representing the task-relevant part; the exogenous reconstruction ô-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:• 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 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. • 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 ô+ but not in ô-, 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 ô-. Owing to space constraints, the complete pre-training results are provided in the appendix. Figure 11 , Figure 12 , Figure 13 , Figure 14 presents the fine-tuned performance of pre-trained models trained with varying pre-training frame durations.Building upon the recurrent state space model (RSSM) architecture [20, 52] , we retain the forward dynamics, posterior encoder, and policy components from the implementation in [52] . The hyper- parameters for SeeX are detailed in Table 2 . We introduce a 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 E θ -mirrors the structure of the endo-encoder E θ + (i.e., the image encoder in [52] ). However, we introduce modifications to the conventional decoder to construct our novel endo-decoder D θ + and exo-decoder D θ -, as expressed in the following:• 1 fully connected layer with 1536 hidden dimensions.• Reshape tensor into shape of (-1, 1536, 1, 1).• 4 transposed convolution layers with 6 output channels, stride 2, and ELU activation with α = 1.0. • Decompose 6-channel input into two 3-channel structures. The first represents the reconstruction output, while the remaining three channels encode the corresponding mask.• 1 output transformation layer of Gaussian distribution.To effectively reconstruct the observed state, we introduce a novel network architecture for the observation model q θ (o t |s + t , s - t ). This network integrates the outputs of the two decoders described earlier to generate the final reconstruction. The detailed network architecture is as follows: • Get two images and two masks from D θ + and D θ -as input.• 1 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. ICM. Curiosity is quantified as the discrepancy between the predicted consequences of the agent's actions and the actual outcomes, as represented in a learned visual feature space. This formulation enables efficient exploration, outperforming baseline methods in VizDoom and Super Mario Bros. Notably, ICM [44] exhibits superior exploration capabilities even in the absence of explicit external rewards. The study underscores the significance of evaluating generalization to new scenarios and offers valuable insights into the effectiveness of reinforcement learning algorithms in adapting to novel environments.Plan2Explore. During exploration, Plan2Explore [52] effectively formulates plans to seek out novel experiences, enabling it to rapidly adapt to downstream tasks in a zero-shot or few-shot setting.Surpassing existing self-supervised exploration techniques, Plan2Explore achieves competitive zeroshot 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 a neural network to predict features of observations extracted from a randomly initialized network, RND [10] effectively enhances exploration in Atari games characterized by sparse rewards. Notably, the method surpasses state-of-the-art performance on challenging games like Montezuma's Revenge, achieving remarkable results without relying on human demonstrations or access to the game'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.LBS. Latent Bayesian Surprise (LBS) [42] leverages Bayesian surprise in a 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 discreteaction environments. It exhibits remarkable resilience to stochasticity in the environment dynamics, facilitating in-depth exploration and achieving superior performance in high-dimensional settings such as video games. Notably, LBS significantly reduces computational costs while enhancing exploration capabilities, making it a promising method for enhancing RL algorithms.APT. APT [38] actively acquires behaviors and representations by exploring novel states in rewardfree environments. This approach maximizes non-parametric entropy in an abstract representation space, circumventing the need for complex density modeling and enabling effective scaling to highdimensional 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-theoretic objective using a 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 a hierarchical agent and ensures skill diversity through maximum entropy policies. Overall, the paper's contributions include proposing a novel method for unsupervised skill learning, showcasing the emergence of diverse skills in various tasks, demonstrating adaptability to new tasks, and providing a foundation for hierarchical reinforcement learning and imitation tasks. DIAYN presents a promising approach for enhancing exploration and data efficiency in reinforcement learning settings.APS. By seamlessly integrating variational successor features with nonparametric entropy maximization, APS [37] effectively optimizes the mutual information between tasks and policy-induced states. This method outperforms existing benchmarks on the Atari 100k data-efficiency 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 lower-dimensional abstract representation space. Empirical results on the Atari benchmark demonstrate state-of-the-art performance, showcasing substantial progress over earlier work. The paper highlights the benefits of utilizing state entropy maximization data for task-conditioned skill discovery and outlines future research directions for further optimization and the integration of diverse approaches. 


Decision: 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.

Reviews: review: paper_summary: paper_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.

They 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.
strength_weakness: strength_weakness: 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
questions: questions: Can you see this method helping in something like a self driving or navigation task in the real world?
limitations: limitations: Yes, they discuss the important limitation that the endogenous and exogenous state are always fully separate and don't interact with one another.
paper_summary: paper_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.

**Post-rebuttal comment:** I believe that the authors have addressed my concerns and I am revising my score from 5 -> 6 and recommend acceptance.
strength_weakness: strength_weakness: 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:
- 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.
- 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.
- 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.
questions: 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: limitations: The authors discuss limitations in L331-334. However, I would like the authors to elaborate on the limitations of their approach.
paper_summary: paper_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.
strength_weakness: strength_weakness: (Not all points are weaknesses. Comments/questions/weaknesses are all gathered in this section for convenience.)

* Line 157: what is BA?
* Line 164: should the second $=$ be $+$?
* 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.
* 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.
* 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.
* Ablation: how does SeeX perform when $\alpha=0$?

Other issues:
* 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.
questions: questions: Please see weaknesses above.
limitations: limitations: The limitations have been briefly discussed in Section 7.
paper_summary: paper_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.
strength_weakness: strength_weakness: - 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.
- 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].
- 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.

[1] Denoised MDPs: Learning World Models Better Than the World Itself

[2] Learning Task Informed Abstractions

[3] Planning to Explore via Self-Supervised World Models

Typos:
- Line 51we use the task-irrelevant (-> task-relevant) information to predict the reward and the action.
questions: 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)?
- 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}$?
- In the experiment, what RL algorithm are the baselines implemented with?
- 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: limitations: N/A

Comments: title: Response
Rebuttal Response
comment: Thank you for your valuable feedback and for the effort you've put into helping us improve our work. We greatly appreciate the insightful discussions we've had with you throughout this process. Rest assured, we will incorporate these discussions and your suggestions into the later formal version of the paper. Your contributions have been instrumental in refining our work, and we are truly grateful for your time and consideration.
Thank you for the reminder. I don't have further questions that require the authors to clarify.
Does the author response address your concerns? Please acknowledge and send any follow up questions.
We hope our responses address your concerns. If you have any further questions or need clarification, please do not hesitate to reach out to us.
Thank you for your kind words and for taking the time to review our responses and discussion. I appreciate your positive feedback and will carefully consider any potential improvements to further enhance the quality of our paper. If you have any additional comments or suggestions in the future, please do not hesitate to reach out.
I thank the authors for their response and discussion! I have no further questions at this time
Thank you for your prompt and thoughtful response to our rebuttal. We appreciate the time and effort you’ve taken to review our rebuttal and provide additional feedback. Below are our responses to the new questions and concerns you raised:
- **Additional ablation results on Walker**
  - As you mentioned in your review, “there are no data points at 200k/300k/400k for baselines,” we understand your concern. To address this, we have included the results in a table, providing more specific quantitative values to better illustrate the performance comparison of different methods at various time steps. Considering your preference for the same type of qualitative line charts as before, we will update the new data points in the line charts of the later formal version.
- **Ablation of $\alpha$​**
  - Firstly, thank you for pointing out in your review that Figure 5 was somewhat confusing. In Figure 5, the top-left, top-right, and bottom-left plots show ablations with different values of $\alpha$ in various environments; since their content is consistent, we present them as bar charts in the PDF for a clearer comparison of their normalized returns. The bottom-right plot depicts ablation experiments of different modules, and we provide a more detailed analysis of these experiments in the table, which shows the specific numerical improvements.
  - Secondly, to address your and future readers' confusion, in the later formal version, we will replace the previous inconsistent Figure 5 with a unified and clearer set of figures. Additionally, we will include the original line plots in the appendix for reference, so readers can still gain a qualitative understanding of the performance of the ablation experiments.
- **Reservations about the novelty**
  - First, thank you for considering that our approach of using only task-related information to measure disagreement is an improvement. Below, I will provide a more detailed explanation of our novelty and contributions.
  - In real-world scenarios, there are many sources of action-independent and task-irrelevant noise. For example, when a drone is patrolling, surrounding buildings, or distant vehicles do not directly affect the drone's flight path or the completion of its patrol task. Many current works overlook this issue, and we are **the first** to address such noise in the URL setting. Additionally, we propose a bi-level optimization framework centered around task-related information and design our methods based on this theoretical foundation. Finally, our experimental results demonstrate the effectiveness of our approach, and ablation studies confirm the utility of different components.
- **Statements without experimental results**
  - Considering the URL problem setting, our model-based framework requires intrinsic reward to assist in early exploration. Intrinsic reward is a modular and can be implemented in various ways. For example, the disagreement measure $Var(\hat{h}^+_{t+1,i})$ used in our work and the variance of the log of predicted probabilities $Var(\log \hat{T}_i(z_t))$ used in [1], both of the intrinsic rewards can be used within our framework.
  - Our main contribution is a theoretically-based framework for solving the URL problem with visual disturbances. In principle, intrinsic reward can be implemented based on the ideas from different methods in URLB[2], as well as the two specific approaches mentioned earlier, meaning our framework is **intrinsic-reward-agnostic**. As you mentioned, the experimental performance of these intrinsic rewards has indeed not been tested. We will consider conducting more extensive and comprehensive comparisons in future work.

Thank you for your continued effort in providing multiple rounds of feedback to improve our work. Your questions have been very valuable. Please feel free to reach out with any further questions or concerns you may have in the future.

[1] Offline Reinforcement Learning from Images with Latent Space Models

[2] URLB: Unsupervised Reinforcement Learning Benchmark
Thank you for reviewing our rebuttal. We are glad to hear that you do not have any new questions. If you have any further concerns or need clarification, please feel free to reach out at any time.
I apologize for my late reply. Thank you to the authors for addressing my questions and providing additional results. I have updated my score to 4, but I still have a few questions:

* Regarding the additional ablation results on Walker, why was a table used instead of a line plot, as in the paper? Additionally, for the ablations related to $\alpha$, why not use a line plot similar to Figure 5 (top-left)? I also didn't understand why different plots were used for the ablation results on $\alpha$ in Jaco and Walker.
* I share similar reservations about the novelty as Reviewer NjAE. To me, the main takeaway from this paper is that using only task-related information to measure disagreement is a better approach. Describing this as a completely new framework is unconvincing. Furthermore, making statements like "Any intrinsic reward can support SeeX's exploration" without experimental results to support it can be misleading and overly generalizing.
Thank the authors for the explanation, I have read the rebuttal and decided to keep my score.
Dear reviewer SZqk,

Since the discussion phase deadline is approaching, we would like to send a friendly reminder.

We greatly appreciate your time and dedication to providing us with your valuable feedback. We hope we have addressed the concerns, but if there is anything else that needs clarification or further discussion, please do not hesitate to let us know.

Thanks, Authors
Dear reviewer NjAE,

Since the discussion phase deadline is approaching, we would like to send a friendly reminder.

We greatly appreciate your time and dedication to providing us with your valuable feedback. We hope we have addressed the concerns, but if there is anything else that needs clarification or further discussion, please do not hesitate to let us know.

Thanks, Authors
Dear reviewer HYxU,

Since the discussion phase deadline is approaching, we would like to send a friendly reminder.

We greatly appreciate your time and dedication to providing us with your valuable feedback. We hope we have addressed the concerns, but if there is anything else that needs clarification or further discussion, please do not hesitate to let us know.

Thanks, Authors
Thank you very much for your thoughtful feedback and for taking the time to review our paper. We sincerely appreciate your constructive comments and are pleased to hear that the new experimental results and clarifications address your main concerns.

We are especially grateful for your recognition of the ablation study involving both branches, and we will certainly incorporate all the feedback provided into the camera-ready version of the paper. Your input has been invaluable in helping us refine our work and ensure its quality.

Thank you once again for your support and for increasing the score. We are committed to making the necessary revisions to meet the expectations and enhance the clarity and impact of our paper.
Thank you for the detailed response to my questions. I believe that the new experimental results and clarification addresses my main concerns, especially the ablation using both branches is helpful in motivating the approach. I have increased my score from 5 -> 6 with the expectation that the authors incorporate all reviewer feedback into a camera-ready version.