Self-Consistent Model-based Adaptation for Visual Reinforcement Learning

1 Introduction

Visual reinforcement learning (VRL) aims to complete complex tasks with high-dimensional observations, which has achieved remarkable results in various domains [Hafner et al., 2019a; Brohan et al., 2023; Li et al., 2024]. Since VRL agents are typically trained on clean observations with minimal distractions, they struggle to handle cluttered observations when deployed in real-world environments with unexpected visual distractions, such as changes in textures or complex backgrounds [Hansen et al., 2020; Fu et al., 2021]. The discrepancy between clean and cluttered observations results in a serious performance gap.

The key to closing the performance gap is to make the policy invariant to distractions. Most existing methods aim to mitigate distractions by learning robust representations. In particular, one line of work is to align the policy's representation between the clean and cluttered observations. Due to the lack of paired data, prevailing methods use hand-crafted functions to create augmentations similar to cluttered observations [Hansen and Wang, 2021; Bertoin et al., 2022]. The effectiveness of such methods is typically limited in settings without prior knowledge of potential distractions. Another line of work addresses the problem through adaptation, which boosts deployment performance by fine-tuning the policy's representation with self-supervised objectives. However, existing adaptation-based methods often lead to narrow empirical increases [Hansen et al., 2020] or are effective only for a specific type of distractions [Yang et al., 2024]. Moreover, the practical application of VRL often requires different policies to ensure robustness against the same types of distractions [Devo et al., 2020]. For instance, domestic robots for different tasks all face the challenge imposed by residential backgrounds with similar distractions. Since policies trained for different tasks have distinct representations, current methods need to fine-tune them separately as the modification made to one policy's representations is not directly applicable to another policy.

∗Corresponding author.

To address the above issues, we propose Self-Consistent Model-based Adaptation (SCMA), a novel method that fosters robust adaptation for various policies as a plug-and-play enhancement. Instead of fine-tuning policies, SCMA utilizes a denoising model to mitigate distractions by transferring cluttered observations to clean ones. Therefore, the denoising model is policy-agnostic and can be seamlessly combined with any policy to boost performance under distractions without modifying its parameters. We further design an unsupervised distribution matching objective to optimize the denoising model in the absence of paired data. Theoretically, we show that the solution set of the unsupervised objective strictly contains the optimal solution in the supervised setting. The proposed objective regularizes the outputs of the denoising model to follow the distribution of observations in clean environments, which we choose to estimate with a pre-trained world model [Hafner et al., 2019b; Hafner et al., 2023].

We practically evaluate the effectiveness of SCMA with the commonly adopted DMControlGB [Hansen et al., 2020; Hansen and Wang, 2021], DMControlView [Yang et al., 2024], and RL-ViGen [Yuan et al., 2024], where the agent needs to complete continuous control tasks in environments with visual distractions. Extensive results show that SCMA

significantly narrows the performance gap caused by various types of distractions, including natural video background, moving camera view, and occlusion. Also, we verify the effectiveness of SCMA with real-world robot data, showing its future potential in real-world deployment. In summary, the main contributions of this paper are:

• We address the challenge of visual distractions by transferring observations and derive an unsupervised distribution matching objective with theoretical analysis.
• We propose self-consistent model-based adaptation (SCMA), a novel method that promotes robust adaptation for different policies in a plug-and-play manner.
• Extensive experiments show that SCMA significantly closes the performance gap caused by various types of distractions. We also demonstrate the effectiveness of SCMA with real-world robot data.

2 RelatedWork

2.1 Visual Generalization in RL

The ability to generalize across environments with unknown distractions is a long-standing challenge for the practical application of reinforcement learning (RL) agents [Chaplot et al., 2020; Tomar et al., 2021; Liu et al., 2023]. Task-induced methods address the problem by learning structured representations that separate task-relevant features from confounding factors [Fu et al., 2021; Pan et al., 2022; Wang et al., 2022]. Augmentation-based methods regularize the representation between augmented images and clean equivalents [Hansen and Wang, 2021; Ha et al., 2023], but they require prior knowledge of the test-time variations to manually design augmentations. Adaptation-based methods [Hansen et al., 2020; Yang et al., 2024] do not assume the distractions and fine-tune the agent's representation through self-supervised objectives. However, existing adaptation-based methods tend to lead to narrow empirical improvement [Hansen et al., 2020] or are limited to a specific type of visual distractions [Yang et al., 2024]. Several studies aim to tackle this issue with foundation models [Nair et al., 2022; Shah et al., 2023], but they still struggle with computational budget and inference time.

2.2 Unsupervised Domain Transfer

Unsupervised Domain Transfer aims to map data collected from the source domain to a related target domain without explicit supervision signals [Wang et al., 2021]. The topic has been explored in various research areas, such as style transfer [Zhu et al., 2017; Zhao et al., 2022], pose transfer [Li et al., 2023], language translation [Lachaux et al., 2020; Artetxe et al., 2017] and so on. However, one key difference between our setting and theirs is that we can interact with the environments to collect data rather than relying on precollected static datasets. Therefore, we can obtain a certain level of control over the distribution of collected data by selecting specific action sequences, which makes it possible for us to achieve the desired transfer from cluttered observations to clean ones with unsupervised distribution matching [Cao et al., 2018; Baktashmotlagh et al., 2016].

3 Methodology

We first present our problem formulation and the supervised objective LO in Sec. 3.1. Then we introduce an unsupervised distribution matching surrogate LKL and analyze the connection between LKL and LO in Sec. 3.2. Finally, we transform LKL into several optimizable adaptation losses in Sec. 3.3, along with practical enhancements in Sec. 3.4.

3.1 ProblemFormulation

We formalize visual RL with distractions as a Noisy PartiallyObserved Markov Decision Process (NPOMDP) Mn = ⟨S,O,A, T ,R, γ, ρ0, fn⟩. In a NPOMDP, S is the hidden state space, O is the discrete observation space, A denotes the action space, T : S ×A ↦→ ∆(S) defines the transition probability distribution over the next state, R : S ×A ↦→ R is the reward function, γ is the discount factor, and ρ0 is the initial state distribution. Here fn : O ↦→ O is a noise function that maps a clean observation ot to its cluttered version ont = fn(ot). Following the common settings [Hansen et al., 2020; Bertoin et al., 2022], we assume that fn is injective so that the distractions do not corrupt the original information. The graphical model of NPOMDP is provided in Fig. 1.

Given the action sequence a1:T , the conditional joint distribution describing the environment's latent dynamics is defined as:

We denote p(ont |ot) = δ(ont - fn(ot)) as the noising distribution of fn, which is a Dirac distribution with δ(·) being the Dirac delta function [Dirac, 1981]. Leveraging the Bayes' rule, the posterior distribution p(ot|ont ) can also be derived from Eq. 1, which we denote as the posterior denoising distribution of fn.

The performance of policies pre-trained with clean observations often degenerates when handling cluttered observations [Hansen et al., 2020; Bertoin et al., 2022]. To fill the performance gap, a natural way is to transfer cluttered observations to their corresponding clean ones by estimating the posterior denoising distribution p(ot|ont ). In the supervised setting, we can estimate p(ot|ont ) with a learnable distribution q(ot|ont ) by minimizing the following objective:

We further show that p(ot|ont ) is a Dirac distribution when fn is injective. Therefore, we adopt a denoising model mde and choose q(ot|ont ) = δ(ot -mde(ont )) in practice. More details can be found in Appendix A.1.

3.2 Mitigating Visual Distractions with Unsupervised Distribution Matching

The direct optimization of LO requires collecting paired observations (ot, ont ). Since we can only collect observations from clean environments (i.e. p(o1:T |a1:T )) and distracting environments (i.e. p(on1:T |a1:T )) separately, the absence of paired data imposes severe challenges. Inspired by unsupervised distribution matching [Baktashmotlagh et al., 2016; Cao et al., 2018], we propose to minimize the KL-divergence between the action-conditioned distribution of the clean and transferred observations, which leads to the following unsupervised surrogate LKL (see Appendix A.2 for details):

where q(o1:T |on1:T ) = ∏t q(ot|ont ) and q(on1:T |o1:T ) = ∏t q(ont |ot) are learnable denoising and noisy distribution respectively.

To analyze the connection between LKL and LO, we first introduce the concept of homogeneous noise functions, which are theoretically indistinguishable in the unsupervised setting and defined below (Details are deferred to Appendix A.2):

Definition 1. For noise functions fni , we denote oni t = fni (ot) as its cluttered observation. Given the distribution of clean observation p(o1:T |a1:T ), we call the noise functions fn1 and fn2 to be homogeneous under p(o1:T |a1:T ) if their cluttered observations have the same distribution, i.e.:

We define Hp fn = {fni |fni ≡p fn}, which includes all homogeneous noise functions of fn under p(o1:T |a1:T ).

Then we show that the solution set of LKL equals the set of posterior denoising distributions of noise functions in Hp fn . Since fn is clearly in Hp fn , the solution set of LKL contains p(ot|ont ), which is the optimal solution to LO.

Theorem 1 (Proof in Appendix A.3). Given p(o1:T |a1:T ) and p(on1:T |a1:T ), let Q denote the solution set of LKL:

It follows that Q equals the set of posterior denoising distributions of noise functions in Hp fn :

Generally speaking, since homogeneous noise functions are theoretically indistinguishable in the unsupervised setting, we can only assure that mde learns to transfer cluttered observations back to clean ones according to a noise function in Hp fn . In Appendix A.3, we further reveal the relationship between the number of homogeneous noise functions and properties of p(o1:T |a1:T ). We also discuss possible ways to reduce the number of homogeneous noise functions in Sec. 3.4 so that Q only contains p(ot|ont ).

To simplify the computation, we show in Appendix A.2 that LKL leads to the following objective, where C is a constant:

Intuitively, the first term regularizes the transferred observations to follow the clean environments' latent dynamics p(o1:T |a1:T ). The second term ensures that the transferred observations remain relevant to the cluttered observations and thus preserve necessary information.

3.3 Adaptation with Pre-trainedWorldModels

Based on the above analyses, we now present the SelfConsistent Model-based Adaptation (SCMA) method, a practical adaptation algorithm that mitigates distractions by optimizing the denoising model with Eq. 3.

Specifically, Eq. 3 involves calculating the actionconditioned distribution p(o1:T |a1:T ), which we estimate with a pre-trained world model [Hafner et al., 2019b; Hafner et al., 2023]. Given a clean trajectory τ = {o1, a1, · · · , oT , aT}, the world model estimates log p(o1:T |a1:T ) with log pwm(o1:T |a1:T ) by maximizing the following evidence lower bound (ELBO):

In the above objective, the KL-divergence objective J t kl enables the model's generation ability by minimizing the distance between the prior and posterior distribution. The reconstruction objective J t o enforces the model to capture the visual essence of the task by predicting the subsequent observations, which facilitates the later adaptation.

Self-consistent Model-based Adaptation. Before adaptation, we first pre-train the policy and world model in clean environments. Then we deploy the pre-trained policy and our denoising model into the distracting environment to collect trajectory {on1 , a1, · · · , onT , aT}, where actions are sampled as at ∼ π(·|mde(ont )). By estimating p(o1:T |a1:T ) with the pre-trained world model, optimizing Eq. 3 leads to the following self-consistent reconstruction loss Ltsc and noisy reconstruction loss Ltn . It should be noted that the policy and world model are both frozen during adaptation. We choose to drop a similar KL-loss term as in Eq. 4 because we empirically find it to have a negative impact on adaptation by harming the reconstruction results, consistent with previous works [Higgins et al., 2017; Chen et al., 2018]. The detailed derivation is provided in Appendix A.4.

Ltsc encourages the denoising model to transfer cluttered observations to clean ones so that the transferred observations will conform to the prediction of the world model. Ltn prevents the denoising model from ignoring the cluttered observations and thus outputting clean yet irrelevant observations. In practice, we implement q(ont |ot) = δ(ont -mn(ot)) with a noisy model mn, and q(o1:T |on1:T ) = ∏t q(ot|ont ) = ∏t δ(ot - mde(ont )) with a denoising model mde.

3.4 Boosting Adaptation by Reducing Homogeneous Noise Functions

As discussed in Theorem 1, the solution set of LKL equals the set of posterior denoising distributions of noise functions in Hp fn . To promote the adaptation, we propose two practical techniques to help the denoising distribution q(ot|ont ) converge to the target posterior denoising distribution p(ot|ont ) by reducing the number of homogeneous noise functions.

Leverage Rewards. If reward signals are available in distracting environments, they can naturally boost adaptation by reducing the number of homogeneous noise functions. Loosely speaking, noise functions with the same p(on1:T |a1:T ) but different p(on1:T , r1:T |a1:T ) are no longer homogeneous if rewards are available. A detailed explanation is provided in Appendix A.3. The derivation in Sec. 3.2 can be simply extended to include rewards by redefining LKL as below (details in Appendix A.4):

which leads to the reward prediction loss:

Ltrew encourages the transferred observations to contain sufficient information of rewards and ignore reward-irrelevant distractions. The final adaptation loss of SCMA is:

Limit the Hypothesis Set of the Denoising Model. For specific types of distractions, we can further encode some inductive bias in the denoising model architecture. Therefore, we can prevent q(ot|ont ) from converging to the posterior denoising distributions of certain homogeneous noise functions by limiting the hypothesis set. For example, we can implement the denoising model as a mask model mmask : Rh×w×c ↦→ [0, 1]h×w×c to handle background distractions. However, to verify the generality of SCMA, we refrain from assuming the type of distractions and implement the denoising model as a generic image-to-image network by default. Detailed implementations are provided in Appendix C.2.

In summary, we propose an adaptation framework with a two-stage pipeline: 1) pre-training the policy and world model in clean environments to master skills and capture the environments' latent dynamics p(o1:T |a1:T ). 2) adapting the policy to visually distracting environments by optimizing q(ot|ont ) with Eq. 5 to transfer cluttered trajectories to clean ones. The pipeline is illustrated with Fig. 2, with further training details provided in Appendix C.

4 Experiment

In this section, we evaluate the capability of SCMA by addressing the following questions:

• Can SCMA fill the performance gap caused by various types of distractions?
• Can SCMA generalize across various tasks or policies from different algorithms?
• How does each loss component contribute to the results? Can SCMA still handle distractions without rewards?
• Can SCMA converge faster compared to other adaptation-based methods or directly training from scratch in visually distracting environments?

(a) video hard

(c) color hard

Table 1: Performance (mean ± std) in visually distracting environments. We report the performance of SCMA and baseline methods in DMControl and RL-ViGen across various distracting settings. The best algorithm is bold for every task.

4.1 Experiment Setup

Environments. To measure the effectiveness of SCMA, we follow the settings from the commonly adopted DMControlGB [Hansen and Wang, 2021; Hansen et al., 2021; Bertoin et al., 2022], DMControlView [Yang et al., 2024], and RL-ViGen [Yuan et al., 2024]. The agent is asked to perform continuous control tasks in visually distracting environments, including video distracting background (video hard), moving camera views (moving view), and randomized colors (color hard). We also evaluate the agent's performance in a more challenging occlusion setting by randomly masking 1/4 of each observation. We provide a visualization of every distracting environment in Fig. 6 in the Appendix. Unless otherwise stated, the result of each task is evaluated over 3 seeds and we report the average performance of the policy in the last episode.

Baselines. We compare SCMA to the state-of-the-art adaptation-based baselines: PAD [Hansen et al., 2020], MoVie [Yang et al., 2024]. We also include comparison with other kinds of methods, including augmentation-based methods: SVEA [Hansen et al., 2021], SGQN [Bertoin et al., 2022], Dr. G [Ha et al., 2023]; and task-induced methods: TIA [Fu et al., 2021], TPC [Nguyen et al., 2021], DreamerPro [Deng et al., 2022]. Following the official design [Hansen and Wang, 2021], the augmentation-based methods use random overlay with images from Place365 [Zhou et al., 2017]. Task-induced methods directly learn the structured representations in distracting environments. Adaptation-based methods will first be pre-trained in the clean environments for 1M timesteps and then adapt to the distracting environments for 0.1M timesteps (0.4M for video hard and 0.5M for RLViGen). By default, SCMA adapts a pre-trained Dreamer policy [Hafner et al., 2019a] to distracting environments. More details can be found in Appendix C.1.

4.2 Adaptability to Visual Distractions

We first evaluate the adaptation ability of SCMA by measuring its performance in the challenging visual generalization benchmarks. Before adapting to the visually distracting environments, we first pre-train the policy and world model in the clean training environment (see Fig. 6a in the Appendix). Then we adapt the agent to visually distracting environments leveraging the pre-trained world model. The experiment results in Table 1 show that SCMA significantly reduces the performance gap caused by distractions and achieves appealing performance compared to augmentation-based methods. While remaining competitive in the color hard setting, SCMA outperforms the best baseline method in most tasks in other 3 settings. Moreover, SCMA obtains the best performance in all tasks in the occlusion setting, which is a common scenario for real-world robot controls. To verify the idea of boosting adaptation by reducing the hypothesis set (Sec. 3.4), we implement the denoising model with specific architectures and conduct experiments in the table-top manipulation tasks with distracting settings from RL-ViGen. Further details are included in Appendix C.2. Following previous work [Yuan et al., 2024], we report the scores under the eval easy and eval extreme settings in Table 1e. The results show that SCMA achieves the best performance in half of the scenarios and remains comparable to other methods in the remaining ones. We believe one way to further improve the performance is to incorporate stronger world models [Ding et al., 2024], which we leave to future works.

We visualize how mde transfers cluttered observations to clean ones in Fig. 3, where it effectively mitigates various types of distractions and restores the task-relevant objects correctly. The qualitative results also indicate that our method can effectively handle distractions not only for large embodiments like walker-walk, but also for challenging small embodiments such as ball in cup-catch and cartpole-swingup, which task-induced methods often fail to manage.

4.3 Versatility of the DenoisingModel

We conduct experiments to measure the versatility of the denoising model from two aspects: 1) can the denoising model generalize across tasks with the same robot? 2) is the denoising model applicable to policies from different algorithms? To answer the above questions, we first cross-evaluate the capability of the denoising model between walker-walk and walker-stand in the video hard environment. Specifically, we take the denoising model adapted to one task and directly evaluate its performance in another task. The results in Table 4 in the Appendix indicate that the achieved denoising model is not restricted to a specific task and exhibits appealing zero-shot generalization capability. To verify that the denoising model is agnostic to policies, we first optimize the denoising model with trajectories collected by a Dreamer policy. Then we combine the obtained denoising model with an SGQN policy in a plug-and-play manner and measure the performance in the occlusion setting. While SGQN reaches appealing results in other settings, it performs poorly under occlusions. However, Table 2 demonstrates that incorporating the denoising model can improve the performance of SGQN by 88%. Therefore, SCMA can serve as a convenient component to promote performance under certain distractions without modifying the policy. However, there is a disparity between the performance of SGQN policy with SCMA and Dreamer policy with SCMA, which we attribute to the policy encoder. As the Dreamer policy's encoder leverages the longterm representation extracted by the world model, it is less susceptible to small mistakes made by the denoising model.

4.4 AdaptationWithout Rewards

While SCMA utilizes both visual and reward signals for the best adaptation results, the ability to adapt without rewards is also important. To address this issue, we conduct experiments in the video hard environments to investigate how different loss components affect the final adaptation results.

To better understand the impact of different losses, we separately removed the 3 loss components from SCMA during adaptation, namely self-consistent reconstruction loss Ltsc , reward prediction loss Ltrew , and noisy reconstruction loss Ltn . From the ablation results in Fig. 15 in the Appendix, we can see that removing the self-consistent reconstruction loss leads to the most significant decrease, indicating that the proposed Ltsc plays a crucial role in adaptation. Another finding is that the reward loss can promote better adaptation by encouraging the denoising model to focus on some miniature yet critical features, such as the ball in ball in cup-catch and the pole in cartpole-balance. While Ltrew contributes considerably to the final adaptation results, Fig. 8 in the Appendix demonstrates that SCMA without rewards still achieves the highest average performance among all adaptation-based methods. The noisy reconstruction loss mainly aims to preserve the connection between the cluttered and transferred observations. Intuitively, removing Ltn from Eq. 3 will cause a mode-seeking problem [Cheng, 1995], where the denoising model will prefer the mode of log p(o1:T |a1:T ) and thus transfer cluttered observations to clean yet irrelevant observations.

4.5 Sample Efficiency in Visually Distracting Environments

Accomplishing tasks with as few cluttered observations as possible is practically important to deploy the agent in distracting environments. Compared with other adaptationbased methods or training from scratch with task-induced methods [Fu et al., 2021; Deng et al., 2022], the performance curves in Fig. 4 show that SCMA can achieve higher performance with much fewer downstream cluttered samples. Although we adapt the policy in video hard with 0.4M steps, SCMA can achieve competitive performance with much fewer steps. We provide the wall clock time and adaptation steps for SCMA to reach 90% of the final performance in Table 5 in the Appendix to show that SCMA obtains compelling results with only 10% of total adaptation timesteps for most tasks.

4.6 Real-world Robot Data

With the rapid development of generative models, their potential to enhance real-world robotic controls has attracted significant attention. Recent works leverage video models to create future observations based on current environment observation and extract executable action sequences with inverse dynamics models (IDM) [Du et al., 2023; Ko et al., 2023]. However, the generated observations might still contain distractions if the input environment observation is cluttered, which imposes challenges to the IDM in making accurate action predictions. We show that SCMA can help IDM better predict the actions when handling cluttered observations. More details are included in Appendix C.3.

We manually collect real-world robot data with a Mobile ALOHA robot by performing an apple-grasping task with teleoperation. The IDM is trained with data collected in the normal setting and evaluated on data collected in 3 distracting settings: 1) fruit bg: various fruits are placed in the background. 2) color bg: the scene is disrupted by a blue light. 3) varying light: the lighting is intentionally changed. We provide the quantitative results in Table 6 in the Appendix and visualization in Fig. 5. The results show that SCMA effectively mitigates real-world distractions and thus has important implications for the practical deployment of robots.

5 Conclusion and Discussion

The ability to generalize across environments with various distractions is a long-standing goal in visual RL. In this work, we formalize the challenge as an unsupervised transferring problem and propose a novel method called self-consistent model-based adaptation (SCMA). SCMA adopts a policyagnostic denoising model to mitigate distractions by transferring cluttered observations into clean ones. To optimize the denoising model in the absence of paired data, we propose an unsupervised distribution matching objective that regularizes the outputs of the denoising model to follow the distribution of clean observations, which can be estimated with a pretrained world model. Experiments in challenging visual generalization benchmarks show that SCMA effectively reduces the performance gap caused by distractions and can boost the performance of various policies in a plug-and-play manner. Moreover, we validate the effectiveness of SCMA with realworld robot data, where SCMA effectively mitigates distractions and promotes better action predictions.

SCMA proposes a general model-based objective for adaptation under distractions, and we wish to further promote this direction by highlighting some limitations and future improvements. SCMA pre-trains world models to estimate the action-conditioned distribution of clean observations. Including stronger world models like diffusion models [Wang et al., 2023] may be a promising way to further promote the performance with complex robots or real-world tasks. Another potential improvement is to explore other types of signals that are invariant between clean and distracting environments, e.g., 3D structures of robots [Driess et al., 2022] or natural language descriptions of tasks [Sumers et al., 2023].

Acknowledgments

This work was supported by NSFC Projects (Nos. 92248303, 92370124, 62350080, 62276149, 62061136001, 92270001), BNRist (BNR2022RC01006), Tsinghua Institute for Guo Qiang, and the High Performance Computing Center, Tsinghua University. J. Zhu was also supported by the XPlorer Prize. Xinning Zhou and Chengyang Ying contributed equally to this work.