Learning Robust Representations with Long-Term Information for Generalization in Visual Reinforcement Learning
Authors: Rui Yang
Jie Wang
Qijie Peng
Ruibo Guo
Guoping Wu
Bin Li
Venue: ICLR 2025 Poster (2025)
Keywords: Visual Reinforcement Learning
Generalization
Representation Learning
Information Bottleneck
One-Step Rewards
Primary area: reinforcement learning
TL;DR: In visual reinforcement learning, we propose a novel robust representation learning approach to simultaneously capture long-term information and discard irrelevant features, facilitating sequential decision-making under unseen visual distractions.

Abstract: Generalization in visual reinforcement learning (VRL) aims to learn agents that can adapt to test environments with unseen visual distractions. Despite advances in robust representations learning, many methods do not take into account the essential downstream task of sequential decision-making. This leads to representations that lack critical long-term information, impairing decision-making abilities in test environments. To tackle this problem, we propose a novel robust action-value representation learning (ROUSER) under the information bottleneck (IB) framework. ROUSER learns robust representations to capture long-term information from the decision-making objective (i.e., action values). Specifically, ROUSER uses IB to encode robust representations by maximizing their mutual information with action values for long-term information, while minimizing mutual information with state-action pairs to discard irrelevant features. As action values are unknown, ROUSER proposes to decompose robust representations of state-action pairs into one-step rewards and robust representations of subsequent pairs. Thus, it can use known rewards to compute the loss for robust representation learning. Moreover, we show that ROUSER accurately estimates action values using learned robust representations, making it applicable to various VRL algorithms. Experiments demonstrate that ROUSER outperforms several state-of-the-art methods in eleven out of twelve tasks, across both unseen background and color distractions.

Published as a conference paper at ICLR 2025
LEARNING ROBUST REPRESENTATIONS WITH LONG- TERM INFORMATION FOR GENERALIZATION IN VISUAL REINFORCEMENT LEARNING
Rui Yangi.,, Jie Wang , Qijie Peng Ruibo Guo", Guoping Wu', Bin Li²
Moo Key Laboratory of Brain-inspired Intelligent Perception and Cognition University of Science and Technology of China
(yr0013, sa24006073, ruiboguo, guoping) @ma uste. edu cn (jiewangx, binli] @ustc. edu.cn
ABSTRACT
Generalization in visual reinforcement learning (VRL) aims to lear agents that can adapt to test environments with unseen visual distractions. Despite advances robust representations learning, many methods do not take into account the essen- tial downstream task of sequential decision-making. This leads to representations that lack critical long-term information, impairing decision-making abilities in test environments. To tackle this problem, we propose a novel robust action-value representation learning (ROUSER) under the information bottleneck (IB) frame- work. ROUSER learns robust representations to capture long-term information from the decision-making objective (i.e., action values). Specifically, ROUSER uses IB to encode robust representations by maximizing their mutual information with action values for long-term information, while minimizing mutual information with state-action pairs to discard irrelevant features. As action values are unknown, ROUSER proposes to decompose robust representations of state-action pairs into one-step rewards and robust representations of subsequent pairs. Thus, it known rewards to compute the loss for robust representation learning. Moreover, we show that ROUSER accurately estimates action values using learned robust representations, making it applicable to various VRL algorithms. Experiments demonstrate that ROUSER outperforms several kute-of-the-ar methods in eleven out of twelve tasks, across both unseen background and color distractions.
INTRODUCTION
Generalization in visual reinforcement learning (VRL) has received considerable attention (Mnih (et al., [2015: Yarats et al., 2021 Kirk al., 2021 Zhu et al. 2023) due to its potential learn agents that can address complex tasks across diverse environments in real-world applications, e.g..
autonomous driving (Zhou al. 2023b Cao al. 2003), robot control (Xing et al. [202T] Liu et al.,
2023b), and chip design (Geng et al. 2024 Wang et al.) 2024cc 2025b). It refers to agents' capability to directly use their learned skills to unknown environments, where visual distractions (e.g.. dynamic backgrounds or object colors under control) may differ from those encountered during training (Li al. 202T Ali et al. 2023). Thus, the generalizable agents can execute tasks with high performance when encountering environments under unseen distractions without extensive retraining.
To learn generalizable agents, one of promising approaches in VRL aims to develop robust repre- sentations against visual distractions in environments (Kemertas & Aumentado-AAmsstrnnn
2021
Mazoure et al.] 2022). Specifically, some methods introduce data augmentations (Hansen & Wang 2021 Huang al. 2022) and contrastive learning (Laskin et al. 2020: Kim et al.) 2021) to learn representations that are robust to irrelevant and spurious information. Other methods apply metric learning (Zhang et al., 2021 Lopez al. 2022) to encode task-relevant information for robust representation learning. Then, based on the learned robust representations, these methods can directly use traditional VRL algorithms (Schulman al. 2017 Fujimoto et al. 2018 HHarmoja al. 2018
"Corresoonding author. Email: jiewangx @ust..cu......

---

Published as a conference paper at ICLR 2025
to optimize the objectives of the downstream task of sequential decision-make i..., maximizing the expected cumulative rewards (i.e., action values) over these environments.
However, the aforementioned representation learning approaches often do not take into account the essential downstream decision-making This results in the representations that cannot capture eritical long-term robust information in sequential data, which is key factor for generalization in VRL (Zhou et al. 2022 Qi et al. 2022 Wang et al., 2023a). Although such representations are robust against visual distractions, they cannot facilitate the generalization of sequential decision-making.
In this paper, we propose a robust action-value representation learning (ROUSER) for VRL general- ization under the information bottleneck (IB). ROUSER introduces IB learn robust representations that capture long-term information from the sequential decision decision-making objective (i.c., action values).
Specifically, ROUSER first applies IB propose a robust action-value representation, which is correlated with the action value. This robust representation for each state-action pair maximizes mutual information with the action value, while minimizing mutual information with the associated state-action pair. Thus, this representation can capture long-term information without irrelevant features. Then, as the true action values are unknown, we cannot directly predict action values to learn the robust action- value representation. Inspired by task decomposition methods (Dietterich) 2000, Russell& & Zimdars, 2003 [ayan 1993 [Brrreto et al.) 2017). ROUSER proposes to decompose the robust action-valle representation into compressed reward pppresentationn--inene captures only information from a one-step reward-ann a subsequent robust action- representation. With this recursive form, ROUSER use the known one-step rewards instead of unknown action values to compute the loss associated with robust action- value representations.
Moreover, for compressed reward representations, ROUSER introduces reward model in the IB framework. This model maximizes mutual information between one-step rewards and compressed reward representations, while minimizing it between such representations and associated state-action pairs. Therefore, ROUSER can use one-step rewards and state-action pairs to learn compressed reward representations that encode only information from rewards.
This study proposes robust action-value representations in the IB framework, capturing long-term information during the robust representation learning process. Thus, it significantly enhances per- formance in sequential decision -eciiion-mmkinggg making against unseen visual distractions. In addition, we provide theoretical guarantee for ROUSER, establishing bound between the true action-value function and the action-value function on top of learned robust action-value representations. This shows that ROUSER can accurately estimate true action values using learned robust action-value representations. Building upon this proof, present an advantage of ROUSER, which is its applicability to various VRL algorithms to improve generalization. We can integrate ROUSER with the VRL critic for action- value estimation, using critic's ' embedding as robust action- value representation. Experiments in Section 5.2 demonstrate the applicability by combining ROUSER with traditional VRL algorithms, including policy gradient (Yarats et al. 2022) and value-based (Hosu & Rebedea 2016) methods.
is worth noting that our study differs from task decomposition methods, which learn the represen- tations from action values but do not take into account robustness against visual distractions. We use one-step rewards to extract long-term information from action values and leverage IB to discard irrelevant features VRL. We summarize the major contributions below.
To the best of our knowledge, this study is the first to learn representations from action values against visual distractions. It encodes robust and long-term information rather than just robust features to facilitate downstream decision-makinn for generalization.
We show that ROUSER action-value representations. With this proof, we can integrate ROUSER with various VRL algorithms to estimate the action values, thereby enhancing robustness.
accurately estimate action values using the learned robust
Extensive experiments demonstrate that ROUSER outperforms several state-of-the--at VRL methods in eleven out of twelve tasks across both unseen background and color distractions.
2
RELATED WORK
Representation Learning for Generalization in VRL. Representation learning approaches for generalization in VRL aim to improve agents' performance unseen test environments by extracting
2

---

Published as a conference paper at ICLR 2025
robust representations from training environments. Some approaches design an auxiliary contrastive task to maximize the mutual information between representations from similar observations while minimizing the mutual information between representations from dissimilar observations (Kim et al.,
2019. Laskin et al., 2020 Fan & Li 2022). To obtain similar observations, they al always leverage data
augmentations generate observations (Hansen et al. [02T]] Haansen & Wang 2021: Huang et al.] 2022, Ye et al. 2023] Wang et al. 2025a). However, they often do not take into account the down-
stream task of sequential decision-making, leading to representations that lack long-term information.
Some approaches introduce the bisimulation metric with a recursive form (Ferns & Precup 2014] Castro 2020 Agarwal et al.) 202T) to identify robust representations using long-term information, similar to our approach. However, it is difficult to directly compute the recursive metric. Thus, for implementation, they often simplify the computation to capture information from the immediate future. Other methods directly predict reward sequences to learn robust representations (Yang et al.) 2022 Zhou et al. 2023a), which captures long-term task-relevant information for generalization. Instead of using reward sequences, approach simply applies one-step rewards to encode the sequential information, and we further introduce the IB principle to improve robustness.
Task Decomposition in RL. We use the task decomposition to represent methods that decompose a task into independent sub-tasks, so that each task has the same environment dynamics but a different reward function. Then, they propose to learn action-valus value representations invariant to different
tasks/reward functions. Some of them, namely successor feature methods (Dayan 1993 Barreto
et al.] 2017 2018: 2020), use such representations to generalize to different tasks for transfer learning.
Others, namely value decompoxition methods (Sutton et al. 1999: Russell & Zimdars 2003 Makino
2023), apply such representations to improve sample efficiecyy in new tasks. See Appendix A for details. Motivated by their framework where action-valle representations can be decomposed into an infinite sequence of reward representations, we (1) learn the action-valus representations for long-term information using one-step rewards and (2) filter out irrelevant information from the one-step reward representations to improve robustness of overall action-value representations.
It is worth noting that in this article, we aim to learn agents in training environments to generalize to test environments with unseen visual distractions, while remaining a consistent reward function. This differs from task decomposition methods, which primarily extract knowledge invariant to tasks between tasks with varying reward functions and no visual distractions. Moreover, task decomposition methods can easily encode irrelevant features in the setting of this article, as they learn representations from environment dynamics, which involve visual distractions in VRL (Yang et al. 2022).
Information Bottleneck (IB). The IB principle in supervised learning (Tishby & Zaslavsky 2015 Saxe et al. 2018) aims to learn compressed representations including minimal information relevant to downstream tasks, improving generalization performance. Specifically, it regularizes representations by minimizing the mutual information between inputs and representations, while maximizing the mutual information between representations and labels. In VRL, recent methods (Pei & Hou 2019 Xiang et al. 2023) also focus on the IB principle that can learn compressed representations to improve generalization capabilities of agents. Different from supervised learning, in VRL, these methods only have one-step rewards but action values, i.c., the targets VRL corresponding labels in supervised learning. Thus, some of them (Federici et al.) 2020 Fan & Li 2022) follow the unsupervised learning paradigm, using the IB framework to regularize representations from similar inputs or multi-view image observations. Since the unsupervised learning in the IB framework is more challenging than supervised learning (Federici et al. 2020), we use one-step rewards-the known supervised signals in RL-to to extract their compressed representations for robust action-value representation learning.
3 PRELIMINARIES
Visual RI. We consider a family of environments E. Each environment E E is a block Markov decision process (BMDP) (Zhang et al. 2020) denoted by M = (S, A, R, p,p", ,7,7) .7) Here, S is the state space, OF is the observation space in e, A is the action space, R is the reward space, p (s',r s. a) is the state transition probability, p" (o',r o, a) is the observation transition probability (i.e.. environment dynamics) varying with the environment
E, and y E [0. 1) is the discount factor. For simplicity, we use bold letters (e.g., o and a) to denote random variables, normal letters (e.g., and a) to denote samples, and O to denote the set of OH for all E E.
3

---

Published as a conference paper at ICLR 2025
We suppose that agent reaches an unseen latent state and obtains observation environment E &. a BMDP, the observation determined by a state and some task-irrelevant visual factors varying with environments, e.g., backgrounds or agent colors in decision- process. Rach state does not involve irrelevant features in VRL and is invariant to the family of environments. Formally, let X he the set of such visual factors. We introduce an observation function g: 8 x X' + O Zhang et al. 2020; Song et al. 2020) such that o = (s, x). Here, x is a random variable in X, which is independent of and with specific transition probability 9f (æ' xx).
Moreover, we assume that the environments follow a generalized Block structure Zhang et al. (2020).
That is, an observation E O uniquely determines its generating state and the visual factor æ. This assumption implies that the observation function g(s, x)) is invertible with respect to both and x. Then, we have pd (o', r 0,a) = r 8, a) qd (x' x:). Note that the expected reward r(o,a) a) = a( Ep [r|s, a] depends on the corresponding state-action pair rather than the observation-action pair, where r is the random variable of the one-step reward.
We aim to learn an agent with policy
O - A(A) that maximizes the expected cumulative reward E 0e 7'r:] simultaneously in all e E E, where [T] means that the expectation taken 1
in e. We use Q (o. ax) = (sc. ar) = Em.p 00 7 ( ac++)] to denote the action values.
ALGORITHM
In this paper, we propose learn robust action-value representations that capture long-term infor- mation for decision-making. Then, we introduce robust and compressed reward representations and use cemporal-difference (TD) learning (Sutton & Barto 2018) paradigm learn robust action-value representations, guided by an IB-based objective to ensure robustness throughout the learning process.
4.1 ROBUST ACTION- VALUE REPRESENTATIONS IN THE IB FRAMEWORK
In this subsection, we propose robust action-value representations in the IB framework.
Firstly, motivated by task decomposition methods, we provide the action-value representation in Definition 4.1 This action-value representation linearly correlated with the action-value function following a given policy, containing the long-term information from the action- value function. Therefore, we propose using the action-value representation to facilitate the agent's decision-making performance in VRL test environments.
Definition 4.1. A representation HT ae) if there exists a linear mapping & Z + IR such that
D-dimensional space Z an action-value representation
Q (ot, 0z) & (HH"(or, oz)) (wr, H (or, at)), Yor V0,E0,u0 az A, where the dimension
1, (+ denotes the inner product, and w denotes the weight vector of $
Then, for VRL generalization, due to visual factors X, the action-value representation may involve irrelevant information. This motivates us to introduce the principle, compressing the action-value representation to discard irrelevant information. Specifically, employ the conditional IB (Chechik & Tishby, 2202) to provide Definition 4.2 this definition, we propose a robust action-value representation, the action-valuc representation for each observation-action pair main- tains maximum mutual information with the corresponding action value while minimizing mutual information with the observation-actinn pair given the action value.
Definition 4.2. An action-value representation is a robust action-value representation 7 (o,, oz) = E [z.] such that for all o O and 0t A,
Z arg min I (u,; H, H.) BI H.), H.
where H, is the random variable of H (o,, at), Q is the random variable of Q" (or, 04), u, is the random variable of (o,, a,).I(y (2) is the mutual information between random variables y and Y2, y ya) is the conditional mutual information to quantify the information between y1 and Y2 given Y3, and B is hyperparameter.
However, it is worth noting that we do not have the true action values during the training process of VRL. and we thus cannot directly use Equations and to learn robust action-value representations.
4

---

Published as a conference paper at ICLR 2025
4.2 COMPRESSED REWARD REPRESENTATIONS FROM ONE- REWARDS
To address the aforementioned challenge posed by unknown action values, we propose using one-step rewards instead. We first present how learn simple action-value representations, and we then discuss how to learn our proposed robust action-value representations in this section.
The Recursive Form of Actien-Value Representations and the Reward Representations. To learn simple action-value representations, following task decomposition methods, we introduce reward representation h Z that is linearly related to the one-step reward, with linear mapping f: 2 + R and the weight vector w'.
a() [h(o,,, a.)) (w', a)))
Based on the reward representations, we can derive the action-value representations such that
HTT (or, 0+)
7 (h(oots;
(or, ae) = (w', (or, az))-
Equation shows that the action-value representation is linearly correlated with the infinite sequence of reward representations. We provide the details Appendix C.T
With Equation can derive recursive form, i.c., the action-value representation for each observation-acee -cction pair consists of a reward representation and the action-value representation of its subsequent observation-actoo -cctionnnationn pair.
Hd (ot, at) = h(o,, az) + +.E., [H (ot+1, a++1)]
(5) simply compute
With the linear relation in Equation and the recursive form in Equation the loss for learning such action-value representations. Moreover, we show that using Equation 5 can converge in tabular settings in Appendix C.2
Compressed Reward Representations for Robust Action-Vable Representations. Based on Definition 4.2 Equations 4 []d as the action-value representations decomposed into infinite sequence of reward representations, which serve as their foundation, we can directly filter out task- irrelevant information from such reward representations improve robustness of overall action-value representations, learning robust action-value representations without irrelevant features. Thus, we propose using the IB framework to regulate the reward representations.
Since (1) the one-step reward is task-relevan without irrelevant features, and (2) the observation may involve irrelevant visual factors, we propose compressed reward representation in Definition 4.3 This representation for each observation- reward and does not involve any irrelevant information by preserving maximum mutual information with this one-step reward while minimizing mutual information with its observation-action pair. Definition 4.3. A reward representation a compressed reward representation z(0e, a,) E [
pair is linearly related to the corresponding one-step
such that for all E O and a E A.
Z, arg arg min I (u, h,(r,) - -I h.),
(6)
h,
where h, is the random variable of h(o, 0t), and is a hyperparameter.
In Equation given the one-step reward r., the conditional mutual information Z(u;; hy/rx) quantifies information h, that is irrelevant. The mutual information I(r;; h:) quantifies task-relevant information shared between the reward r, and the representations h, 3 is hyperparameter that determines the preference over the trade-off between lask-relevant and task-irrelevant information.
Based on Definition 4.3 and Equation 4 the robust action-value representation is linearly correlated with the infinite sequence of compressed reward representations. Thus, we can also use the recursive form of robust action-value representations, as shown in Equation 5 to learn these representations.
Furthermore, in the following of this part, we discuss how to compute mutual information on Equation 5 We provide detail formulas of Equation 6 giveen the joint distribution and conditional joint distribution of u, Zt, and r,. Note that for simplicity, we use instead of h.
p(2: (u.) log
p(re|z))
I(u 24 2r/ra) Eu
I(re: z4) Erazs log
p(z. r.)
p(r.)
5

---

Published as a conference paper at ICLR 2025
To estimate the probability density p(z, (u) on Equation 7 assume that compressed reward representation 2(0,, ar) for any abservation-actiin pair (o,, ar) follows a Gaussian distribution Ne., and we provide the detailed implementation in the next subsection.
However, due to the modelling of zy the approximation of p(z, 12) is intractable. Inspired by the method that estimates conditional mutual information (Alemi al. 2017 Fan & Li 2022). we intro-
duce a new representation k, E 2 extracted directly from the one-step reward F, to estimate p(k,|r,) instead of p(z)(rr). Thus, we can derive an upper bound of the conditional mutual information.
I(u; (r:))) S Drcc (piza (l))lll
(8)
We provide details in Appendix C.3 We then can minimize the upper bound-i.e., the Kullback- Leibler (KI.) divergence between p(z; (t)) and p(k,/rr)-tto discard irrelevant and redundant infor- mation for compressed reward representation learning. Moreover, the representation kę is directly mapped from the one-step reward rt. To prevent this mapping from introducing potential redun- dant information into k,, we first use I(ry; ke/ux) ) to quantify this information and then minimize T(ry; kg |.). Therefore, we have the whole optimization objective to minimize for compressed reward representation learning.
Lebb I(u:; - -I((r; (.
4.3 ROUSER ALGORITHM
Here, we present the detailed algorithm of ROUSER. Firstly, we provide the losses for compressed reward representations. Then, introduce the loss with TD learning paradigm for robust oetion-value representation learning. Moreover, we provide the theoretical analysis of learned robust action-valus representations. We provide the architecture and pseudocode of ROUSER in Appendix D
IB-based Losses for Compressed Reward Representation Learning. The optimization objective for learning compressed reward representations consists of two parts: prediction loss based on Equation and conditional IB loss based on Equation 9
Firstly, we introduce reward model & encode the compressed reward representation from an observation-action pair and assume that compressed reward representation =(0r, a.) = B[z,] for any observation-action pair (Or, a,) follows a Gaussian distribution No N(z (ot, a,), (or, 4:)). where zy (0t, ae) is the mean and (o,, 0²) the diagonal covariance. We also provide a parameter- ized model 1) estimate the representation k(r.) = E. k. k, is drawn from the Gaussian distribution Nr = N (k, (r:)). where (rx) is the mean of k,, and En (r) is the diagonal covariance.
Secondly, we provide the prediction loss based on Equation 3 to maintain the linear mapping relationship between the compressed reward representations and corresponding one-step rewards. Building upon model-based RL methods (Janner et al. 2019, Wang et al. [022 2023c) that model one-step rewards as Gaussian distribution, predict the mean and the standard deviation of reward distribution to learn reward representations. We formulate the prediction loss based on KI. divergence:
(WW,==((0r at)) rall2 + log (w", (0 an)) a:))
(10)
2 (w", zw az))²
where w' denote the weight vector mentioned in Equation for the computation of the mean of one-step rewards, w" denote the weight vector to compute the standard deviation of one-step rewards. and B denotes the replay buffer.
Finally, we propose the conditional IB loss based on Equation to learn compressed reward represen- tations, discarding task-irrelevant information while preserving task-relevant information. With the models ya and 7), we can reformulate Equation as:
Eg I(z;;) + DDece (MoIIA.))
Linto(d:: (bb,o 17)
(11)
where derived from 3 is a hyperparameter that is tuned by exponential scheduler during training, and DsKc represents the symmetrized KL divergence (Jeffreys] 1998) obtained by averaging the KL divergences Dx.lWNll),, and Dxc(N, (. Please refer Appendix C.3

---

Published as a conference paper at ICLR 2025
The first term of Linto in Equation 11 maximizes mutual information between the representation z, and the one-step reward rt. This term in fact maximizes task-relevant information in the representation ze. The second term of Lints minimizes the conditional mutual information between the observation- action pair (Or, a,) and the representation 2; It represents that we can minimize irrelevant and redundant information in the representation 24- See Appendix C.3 for the derivation of Equation 11
Further, to estimate the mutual information I(z,; k,) between 2, and k, for a transition (o, a, r) ~ B, we can leverage the InfoNCE (van den Oord et 2018 Laskin
2020) function. Thus, we have
exp(sim(zz; kz))
I(ze; kt)
log exp(sim(Zz: kg)) + > i=1 exp (sim (Ze, kt,;))
(12)
where si = x² y, N is the batch size.
TD-Based Loss for Robust Aetion-Value Representation Learning. As the robust action- value representation has recursive nature, we employ the TD learning paradigm the commonly used technique to learn an objective with a recursive form-to learn the representation.
Based on Equations and 6 we derive recursive form of the robust action-value representation. 7" (or, ax) = 2 2(,, a.) + y|E=T [Z" a+41)]
(13)
where 2d and subject to Equations] and respectively.
Then, we model the critic using the parameter 0 to estimate the action values. For robust action-value representation learning, we do not change critic's structure and directly use embeddings from critic's center layer as robust action-value representations (see Figure I. Therefore, we also parameterize the robust action-value representations as Zg. Based on Equation 13 we have the TD-based loss as:
Crobust (0) = Es ||| Zo (o,, ae) - Z(o, a )1.]] [o, a²) (o,, a.) 7Exx [Zg (o a
(14)
where (o:, ac, rz,; 0r+1) B, and 0 denotes the parameter of the target critic that are obtained as an exponentially moving average of the parameter 0 of the critic (Mnih et al. 2015). Moreover, if we apply the soft actor-critic framework (Haarnoja al.) 2018) that adds entropy the policy to one-step rewards, we will also add this entropy into each dimension of compressed reward representations.
According to Equation 14 we provide a theoretical guarantee for the robust action-value representa- tion learning in the following theorem. It gives bound between the true action-value function and the action-value function on top of learned robust action-valus representations (see Appendix C.4). Theorem 4.4. Let be a learned robust action-value representation from any observation-action pair (o,a) O x A, Q O x A + R be the true action-value function of a policy in the environment E E, ft Z + R be the optimal linear mapping on the representation space, and e be a bound of estimation ervor for each compressed reward representation z, i.e., lf)(2) (2)--([l) r| For
any (o, a) E O x A and E E, we have
1 0 s 12 (2 T-7
This theoretical proof allows us to directly use embeddings of the critic's center layer as robust action-value representations for estimating action values. This shows that ROUSER can be integrated with various VRL methods. We combine ROUSER with traditional VRL methods in Section 5.2
In addition, alleviate the overestimation bias of action values, Fujimoto et al. (2018) introduces double critics with 01 and fz. We note that our representation learning based on Equation 14 also suf- fers from the overestimation bias well as the actie function approximation (Fujimoto et al.) 2018). Thus, we propose to leverage the overestimation bias of robust action-value representations to alleviate the bias of estimated action values. We reformulate Equation 14 as:
Crobust (0, (0, = ((
Σ 1=1.2
Ze, (o,, a.) 2(0,, ar)
(15)
Ž(0,, ar) = 20 (o,, ax) + TLE
minl Lg, (0++1, ac+1) 1=1,2
minl Zg, (or, a+) min T0 Zg, (o,, ar) + max Tureo (or, ar), i1,
(16)
==1
11,
7

---

Published as a conference paper at ICLR 2025
Actor
Actur Loss
02
Encoder
ROUSER
0, 02
dp
Q
Critic Loss
a22
Z
Robust Loss
robust
Reward Modkl
Vl, 0aaina
Reward
z
(1re w wo',""" Linto (1). r1)
n
k
Figure 1: The architecture of ROUSER. We first input (or, Qz, Tx) into the reward model to compute the reward loss, updating the reward model. Next, we input (0r, ar) into the critic model to compute the robust loss, which updates (01, 02) to learn robust action-valle value representations. Morcover, use traditional actor and critic losses to update the actor and the entire critic model, respectively.
where 0, and 02 are parameters of target critics, Iu-20 and Iu'co denote D-dimensional vectors of binary values from the indicator function I, and w/ is the parameter of the linear mapping f equivalent D. With the operator minI, we select minimums robust action-value representations positively related to action values and maximums negatively related to action values. Thus, based on the linear relation between action values and robust action-valus representations, we can alleviate the overestimation bias in the action-value estimation.
To clearly overview our method, we provide the architecture Figure I illustrates our proposed
two loss functions: the robust loss Loousa (0,, 02) and the reward loss Eprea w/, wf,w") + Linfo (nn 1).. The detailed implementation of this architecture are presented in Appendix
EXPERIMENTS
We conduct extensive experiments to evaluate the generalization performance of ROUSER on testing environments with unseen visual distractions. Firstly, we demonstrate the effectiveness of ROUSER for solving the visual control tasks on DeepMind Control Suite (DMC. Tassa et al. (2018) with task-irrelevant visual distractions, including background and color distractions. Then, we assess the applicability of ROUSER by combining it with traditional policy gradient and value-based VRL algorithms. Finally, we visualize the captured long-term information and evaluate the robustness of ROUSER. Moreover, we conduct careful ablation studies to show the effectiveness of ROUSER. We provide additional results of ROUSER in Appendix E
5.1 GENERALIZATION PERFORMANCE ON DMC BENCHMARK WITH CONTINUOUS ACTIONS
We apply visual distractions in DMC benchmark, including color and background distractions (Stone) al.) 2202T) to investigate the generalization performance of VRL agents. For color distractions, we train the agents in training environment with weak color change and evaluate in various test environments with strong color changes. It is worth noting that the variance the color distractions we use is larger than the hard setting proposed in Stone et al. (2021), which poses greater challenges for the agents. For background distractions, we follow the hard setting (Agarwal al. 202T Yang et al., (022) to use 2 videos as backgrounds during training and evaluate the generalization performance on 30 unseen videos. See Appendix [E.1 for more details. We conduct main experiments on six tasks of DMC. For each task, train the agents with six random seeds and present means and standard errors of their cumulative rewards at 500,000 environment steps.
We integrate ROUSER with DrQv2 (Yarats et al. 2022 and SRM (Huang et al. 2022 evaluate the effectiveness of ROUSER on generalization. Note that in all experiments under both color and background distractions, ROUSER-DrQv2 consistently outperforms DrQv2, and ROUSER-SRM
8

---

Published as a conference paper at ICLR 2025
Color ball in cup catch
Color finger spin
Video ball in catch
14 Video finger spin
Cumulative Reward
Comulative Reward
Cumulative Reward
Comulative Reward
Cumulative Reward
-5
Color cartpole swingup
Cumulative Reward
Color walker walk
Cumulaative Reward us -3 32
cartpole swingup
Cumulative Reward 2- se
Video walker
Color Chletach
Color reacher easy
Video cheetah
i2. Video reacher easy
Cumul Cimulative Reward
Cumulative Reward
Cummulati Inmulative Reward
Cunulative Reward
2
ROUSER-SRM
CRESP
ROUSER-DIQV2
CURL
SRM
RAP
DrQv2
TACO
Figure We evaluate ROUSER on DMC with unseen color and video background distractions. Each result is averaged over 100 episodes at 500K environment steps using six random seeds.
surpasses SRM. We also compare approach ROUSER against several Mfate-of-the-ee (SOTA)
methods: (1) CURL (Laskin et al. 2020), (2) RAP (Chen & Pan (022). (3) CRESP YYang et al.
2022), (4) TACO (Zheng al. 2024). See Appendix E. for detailed settings.
Figure illustrates results under color and background distractions, demonstrating that ROUSER outperforms other baselines in 11 out of 12 experiments. See Appendix H.T for detailed results.
is worth noting that although ROUSER uses one-step rewards to learn robust action-value represen- tations, it significantly improves generalization performance on tasks with sparse reward functions finger_spin, and reacher_ean easy tasks from DMC benchmark). Considering
(e.g., that the actor-critic algorithms (Fujimoto al. 2018 Haarnoja et al. 2018 which estimate ac- tion values using one-step rewards also perform well on sparse reward tasks, we believe that our approach ROUSER also effective the contexts beyond sparse rewards.
5.2 EXTENDING ROUSER TO PROCGEN WITH DISCRETE ACTIONS
In this subsection, motivated by the result in Theorem 4.4 extend ROUSER decision tasks with discrete actions, assessing its applicability.
ROUSEU QRIINNN
17.5
QU-DEN
15.0
12.5
Firstly, Figure 2 shows that ROUSER can enhance general- ization when combined with policy gradient VRL algorithms, such DrQv2 and SRM. Then, demonstrate that ROUSER can be integrated with value-based methods, we conduct ex- periments Procgen (Cobbe et al. 2020). benchmark with image observations where each game has multiple levels to evaluate generalization. See Appendix F.2 for detailed settings.
10.0
7.5
10 Sooro
5.0
2.5
0.0
dodgeball cavetlyer
bigfish
Figure 3: Results on Procgen.
We integrate ROUSER with QR-DQN (Dabney et al.) 2018), prior SOTA value-based VRL method on Procgen for generalization.
Figure 3 shows Procgen scores on test levels that averaged over five random seeds with batch size of 256 at 25M environment steps. The results demonstrate that ROUSER can also improve the generalization performance of agents when combined with the value-basse VRL method for discrete control tasks. In addition, we provide the details of these results and the implementation for the combination of ROUSER and QR-DQN in Appendix E.2

---

Published as a conference paper at ICLR 2025
1e
162
ROUSER IIQQ2
4
1e2
BOUUERR-DVQY
Conmulatirr Rewarrds
le2
MSE Loss
ROLSER-SMM
ExEB
1.
D5:2
1216
DIQV2
Commldiie
Curmulativn
DS9
DIS
0
Epochs
lel
Environment Steps
le5
Enviroament Steps
3
les
at 500K Environment Sleps
(a) Long-Tere -Tem Information
(b) Conditional
(c) Noise in Rewards
(d) Different Dimensions
Figure 4: We report four results on cheetah run task with unseen color distractions.
5.3 ANALYZING THE REPRESENTATION LEARNED BY ROUSER
Long-Term Robust Information. We assess how ROUSER can capture long-term robust information. Specifically, first collect 200K samples under unseen color distractions using DrQv2's policy. Next, we compute the representations learned by each method from state-action pairs in the collected data. Such representations are extracted from the center layer of the critic in each method. Then, we fix these representations and input them into a 2- 2-layer trainable MI.P to predict the future reward sequences over a length of 300. During the prediction process, we use the Adam optimizer with a learning rate le-4. As shown Figure 4 a), we plot curves of the MSE loss IK evaluation samples across three seeds. The lower MSE losses ROUSER indicate that its learned representations capture long- term information more effectively.
649.7
488.0
363.0
316.1
663.7
524.4
398.4
283.3
683.1
510.3
367.9
251.4
0.3
0.4
0.5
0.6
Distraction Scale
Figure 5: Results with differ- ent a in different color scales.
Effectiveness of IB. To assess the effectiveness of IB, we remove the conditional IB loss from the objectives of ROUSER. We then train agents to evaluate their generalization performance. Results (i.e., means and standard deviations) over three seeds in Figure (b) demonstrate the effectiveness of the conditional IB loss. Moreover, by tuning that larger values of a result in smaller DsKe, potentially improving generalization performance. This observation aligns with our analysis Section 4.3, demonstrating that minimizing Dskc (the second term of Linfo) effectively discards irrelevant and redundant information. We illustrate the results for different values a over three seeds in Figure 5 Each result is evaluated at 200K steps on cartpole-swingup task under color distractions (scale from 0.3 to 0.6, see Appendix E.5 details).
the conditional IB loss from Equation [T] we find
Noisy One-Step Rewards. We add noise drawn from (0, 0.1) for each reward to show how
ROUSER be robust against noisy rewards. Figure 4c) demonstrates that ROUSER-DrQv2 significantly outperforms DrQv2, remaining effective under the noisy reward setting.
Ablation Study for Dimension. We illustrate results ROUSER-DiQV2 for different dimensions of robust action-value representations in Figure dd. The results, averaged over three seeds on cheetah_ run task with color distractions, show that ROUSER is not sensitive to the dimension D.
CONCLUSION
Many existing VRL approaches focus on learning robust representations against irrelevant visual distractions but often do not take into account the long-term decision-relevant information. This may degrade generalization performance of VRL agents. In this paper, we propose robust action-value representations in the IB framework, which encodes long-term information from action values while discarding irrelevant features from image observations during the robust representation learning process. Considering that action values unknown, we propose decompose robust action-value representations into infinite sequences associated with rewards. Thus, we can directly use known rewards to learn robust action-value representations. Experiments demonstrate that our approach significantly enhances generalization capabilities in the majority of tasks. We believe this method is general and applicable other real-wor wellworr applications, such end-to-end autonomous driving.
10

---

Published as a conference paper at ICLR 2025
REPRODUCIBILITY STATEMENT
In this study, to ensure the reproducibility of our approach, we provide key information from the main text and Appendix as follows.
1. Algorithm.
We provide the architecture and pseudocode of our approach ROUSER in We also provide the detailed implementation of ROUSER in Appendices E.] E.5 for the hyperparameters of ROUSER.
Appendix D
and
E.2 See Appendix
2. Source Code. According to the architecture and pseudocode in Appendix our approach ROUSER can directly build upon traditional VRI. algorithms. Specifically, Section 5.1 we follow Yuan et af.] (2023) to use the provided code for DrQv2 and SRM, available at In Section 5.2 apply the code from Wang et al. (2020) for QR-DQN, available at httpp://gittubb.cmm http://gtuuub.ooomaaa www/kaiiin9//miiree.... Moreover, we implement our codes in Python version 3.7 and make the code available online
3. Experimental Details. We provide detailed experiment settings in Section 5.1 Appen-
dices E.1 and
E.3
4. Theoretical Proofs. We provide all proofs in Appendix
ACKNOWLEDGMENTS
The authors would like to thank all the anonymous reviewers for their valuable suggestions. This work supported by the National Key R&D Program of China under contract 2022ZD01 19801 and the National Nature Science Foundations of China grants U23A20388 and 62021001.
REFERENCES
Majid Abdolshah, Hung Le, Thommen George Karimpanal, Sunil Gupta, Santu Rana, and Svetha Venkatesh. A new representation of successor features for transfer across dissimilar environ- ments. Proceedings of 38th International Conference on Machine Learning, volume 139 of Proceedings of Machine Learning Research, pp. 1-9. PMLR, 2021.
Rishabh Agarwal, Marlos C. Machado, Pablo Samuel Castro, and Mare G. Bellemare. Contrastive behavioral similarity embeddings for generalization in reinforcement learning. In 9th International Conference on Learning Representations, 2021.
Alexander A. Alemi, Ian Fischer, Joshua V. Dillon, and Kevin Murphy, Deep variational information bottleneck. In 5th International Conference Learning Representations, 2017.
Haider Ali, Dian Chen, Matthew Harrington, Nathaniel Salazar, Mohannad Al Ameedi, Ahmad Faraz Khan, Ali Raza Butt, and Jin-Hee Cho. A survey on attacks and their countermeasures in deep learning: Applications in deep neural networks, federated, transfer, and deep reinforcement learning. IEEE Access, 11:120095-120133 2023.
André Barreto, Will Dabney, Rémi Munos, Jonathan J. Hunt, Tom Schaul, David Silver, and Hado van Hasselt. Successor features for transfer in reinforcement learning. In Advances in Neural Information Processing Systems 30, pp. 4055-4065, 2017.
André Barreto, Diana Borsa, John Quan, Tom Schaul, David Silver, Matteo Hessel, Daniel J. Mankowitz, Augustin Zidek, and Rémi Munos. Transfer in deep reinforcement learning using successor features and generalised policy improvement. In Proceedings of the 35th International Conference on Machine Learning, volume 80 Proceedings of Machine Learning Research, pp. 510-519. PMLR, 2018
André Barreto, Shaobo Hou, Diana Borsa, David Silver, and Doina Precup. Fast reinforcement learning with generalized policy updates. Proc. Natl. Acad. Sci. USA, 117748):30079--30057 2020.

---

Published as a conference paper at ICLR 2025
David Bertoin, Adil Zouitine, Mehdi Zouitine, and Emmanuel Rachelson. Look where you look! crriency-guidde q q-networks for generalization visual reinforcement learning. In Sanmi Koyejo. S. Mohamed, A. Agarwal, Danielle Belgrave, K. Cho, and A. Oh (eds.), Advances in Neural Information Processing Systems 35, 2022.
Zhong Cao, Kun Jiang, Weitao Zhou, Shaobing Xu, Huei Peng, and Diange Yang. Continuous improvement of self-driving cars using dynamic confidence-aware reinforcement learning, Nature Machine Intelligence, 5(2):145-155. 2023.
Wilka Carvalho, Andre Saraiva, Angelos Filos, Andrew K. Lampinen, Loic Matthey, Richard L. Lewis, Honglak Lee, Satinder Singh, Danilo Jimenez Rezende, and Daniel Zoran. Combining behaviors with the successor features keyboard. In Advances in Neural Information Processing Systems 36, 2023.
Pablo Samuel Castro. Scalable methods for computing state similarity in deterministic markov decision processes. In The Thirty-Fourth AAAI Conference on Artificial Intelligence, pp. 10069- 10076. AAAI Press, 2020.
Gal Chechik and Naftali Tishby. Extracting relevant structures with side information. In Advances in Neural Information Processing Systems 15, pp. 857-864, 2002.
Jianda Chen and Sinno Jialin Pan. Learning representations via robust behavioral metric for deep reinforcement learning. In Sanmi Koyejo, S. Mohamed, A. Agarwal, Danielle Belgrave, K. Cho, and A. Oh (eds.), Advances in Neural Information Processing Systems 35, 2022.
Karl Cobbe, Christopher Hesse, Jacob Hilton, and John Schulman. Leveraging procedural generation to benchmark reinforcement learning. In Proceedings of the 37th International Conference on Machine Learning, volume 119 of Proceedings of Machine Learning Research, pp. 2056.2056.
PMLR, 2020.
Will Dabney, Mark Rowland, Mare G. Bellemare, and Rémi Munos. Distributional reinforcement learning with quantile regression. In Sheila A. Mcllraith and Kilian Q. Weinberger (eds.), Proceed- ings of the Thirty-Second AAAI Conference on Artificial Intelligence, pp. 2892-2901. AAAI Press, 2018.
Peter Dayan. Improving generalization for temporal difference learning: The successor representation. Neural Comput., 5(4):613-624 1993.
Thomas G. Dietterich. Hierarchical reinforcement learning with the MAXQ value function decompo- sition. J. Artif. Intell. Res., 13:227-303, 2000.
Jiameng Fan and Wenchao Li. DRIBO: robust deep reinforcement learning via multi-view information bottleneck. In International Conference Machine Learning, volume 162 of Proceedings of Machine Learning Research, pp. 6074-6102. PMLR, 2022.
Linxi Fan, Guanzhi Wang, De-An Huang, Zhiding Yu, Li Fei-Fei, Yuke Zhu, and Animashree Anandkumar. SECANT: self-expert cloning for zero-shot generalization of visual policies. In Marina Meila and Tong Zhang (eds.), Proceedings of the 38th International Conference on Machine Learning, volume 139 of Proceedings of Machine Learning Research, pp. 3088-3099, PMLR, 2021.
Mehdi Fatemi and Arash Tavakoli. Orchestrated value mapping for reinforcement learning. In The Tenth International Conference on Learning Representations. OpenReview.net, 2022.
Marco Federici, Anjan Dutta, Patrick Forré, Nate Kushman, and Zeynep Akata. Learning robust
representations via multi-view information bottleneck. In 8th International Conference Learning Representations. OpenReview.net, 2020.
Norman Ferns and Doina Precup. Bisimulation metrics are optimal value functions. In Proceedings of the Thirtieth Conference on Uncertainty in Artificial Intelligence, pp. 210-219. AUAI Press, 2014.
12

---

Published as a conference paper at ICLR 2025
Scott Fujimoto. Herke van Hoof, and David Meger. Addressing function approximation error in actor-critic methods. In Proceedings of the 35th International Conference on Machine Learning, volume 80 of Proceedings of Machine Learning Research, pp. 1582-1591. PMLR, 2018
Zijie Geng, Jie Wang, Ziyan Liu, Siyuan Xu, Zhentao Tang, Mingxuan Yuan, Jianye HAO, Yongdong Zhang, and Feng Wu. Reinforcement learning within tree search for fast macro placement. In Forty-first International Conference on Machine Learning, 2024. URL /ttp:ttttp/////eeeeeeeee net/forurd det.:orrmmmi=--.i.. A=JGMSXORRUV
Christopher Grimm and Satinder Singh. Learning independently-obuiinbbee reward functions. CoRR, abs/1901. abs/1901.08649, 2019.
Jeongsoo Ha, Kyungsoo Kim, and Yusung Kim. Dream to generalize: Zero-shot model-based reinforcement learning for unseen visual distractions. In Brian Williams, Yiling Chen, and Jennifer Neville (eds.), Thirty-Seventh AAAI Conference on Artificial Intelligence, pp. 7802-7810. AAAI Press, 2023.
Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with stochastic actor. In Proceedings of the 35th International Conference on Machine Learning, volume 80 of Proceedings of Machine Learning Research, pp. 1856-1865. PMLR, 2018.
Nicklas Hansen and Xiaolong Wang. Generalization in reinforcement learning by soft data augmen- tation. In IEEE International Conference on Robotics and Automation, pp. 13611-13617. IEEE, 2021.
Nicklas Hansen, Hao Su, and Xiaolong Wang. Stabilizing deep q-learning with convnets and vision transformers under data augmentation. In Advances in Neural Information Processing Systems 34, pp. 3680-3693, 2021.
Matteo Hessel, Joseph Modayil, Hado van Hasselt, Tom Schaul, Georg Ostrovski, Will Dabney, Dan Horgan, Bilal Piot, Mohammad Gheshlaghi Azar, and David Silver. Rainbow: Combining im- provements in deep reinforcement learning. In Proceedings of the Thirty-Second AAAI Conference on Artificial Intelligence, pp. 3215-3222. AAAI Press, 2018
Ionel-Alexandru Hosu and Traian Rebedea. Playing atari games with deep reinforcement learning and human checkpoint replay. CoRR, 0ans//607.05000 2016.
Yangru Huang, Peixi Peng, Yifan Zhao, Guangyao Chen, and Yonghong Tian. Spectrum random masking for generalization in image-based reinforcement learning. In Advances Neural Infor- mation Processing Systems 35, 2022.
Michael Janner, Justin Fu, Marvin Zhang, and Sergey Levine. When trust your model: Model-based policy optimization. In Advances in Neural Information Processing Systems 32, pp. 12498-12509, 2019.
Harold Jeffreys. The theory of probability. OuP Oxford, 1998.
Mete Kemertas and Tristan Aumentado-Armstrone Towards robust bisimulation metric learning. In Advances in Neural Information Processing Systems 34, pp. 4764-4777, 2021.
Hyoungseok Kim, Jaekyeom Kim, Yeonwoo Jeong, Sergey Levine, and Hyun Oh Song. EMI: exploration with mutual information. In Proceedings of the 36th International Conference on Machine Learning, ICML 2019, 9-15 June 2019, Long Beach, California, USA, volume 97 of Proceedings of Machine Learning Research, pp. 3360-3369. PMLR, 2019.
Jackyeom Kim, Minjung Kim, Dongyeon Woo, and Gunhee Kim. trop-bottleneck: Learning discrete compressed representation for noise-robust exploration. In 9th International Conference Learning Representations. OpenReview.net, 2021.
Robert Kirk, Amy Zhang, Edward Grefenstette, and Tim Rocktäschel. survey of generalisation in deep reinforcement learning. CoRR, abs/2111.09979 2021.
13

---

Published as a conference paper at ICLR 2025
Yufei Kuang, Jie Wang, Haoyang Liu, Fangzhou Zhu, Xijun Li, Jia Zeng, HAO Jianye, Bin Li, and Feng Wu. Rethinking branching on exact combinatorial optimization solver: The first deep symbolic discovery framework. In The Twelfth International Conference Learning Representations, 2024.
Michael Laskin, Aravind Srinivas, and Pieter Abbeel. CURL: contrastive unsupervised represetttttssss for reinforcement learning. In Proceedings of the 37th International Conference on Machine Learning, volume 119 of Proceedings of Machine Learning Research, pp. 5639-5650. PMLR, 2020.
Bonnie Li, Vincent François-Lavet, Thang Doan, and Joelle Pineau, Domain adversarial reinforcement learning. CoRR, abs/2102.07097, 2021.
Timothy P. Lillicrap, Jonathan J. Hunt, Alexander Pritzel, Nicolas Heess, Tom Hrez, Yuval Tassa, David Silver, and Daan Wierstra. Continuous control with deep reinforcement learning. In 4th International Conference on Learning Representations, 2016.
Zichuan Lin, Li Zhao, Derek Yang, Tao Qin, Tie-Yan Yan Liu, and Guangwen Yang. Distributional reward decomposition for reinforcement learning. In Advances Neural Information Processing Systems 32, pp. 6212-6221, 2019.
Haotian Ling, Zhihai Wang, and Jie Wang. Learning to stop cut generation for efficient mixed- integer linear programming. Proceedings of the AAAI Conference on Artificial Intelligence, 38(18): 20759-20767, Mar. 2024.
Haoyang Liu, Yufei Kuang, Jie Wang, Xijun Li, Yongdong Zhang, and Feng Wu. Promoting generalization for exact solvers via adversarial instance augmentation, 2023a. URL https:
7/arxiv. org/ abs /atos///110.mmm 14161
Haoyang Liu, Jie Wang, Wanbo Zhang, Zijie Geng, Yufci Kuang, Xijun Li, Bin Li, Yongdong Zhang, and Feng Wu. MILP-studio: MILP instance generation via block structure decomposition. In The Tiirth- Annual Conference on Neural Information Processing Systems, 2024. URL
Haoyang Liu, Jie Wang. Zijie Geng, Xijun Li, Yuxuan Zong, Fangzhou Zhu, Jianye HAO, and Feng Wu. Apollo-MILP: An alternating predictioncooreccton neural solving framework for mixed-integer linear programming. In The Thirteenth International Conference on Learning Representations, 2025. URL pttttpp:///penrevieey net/ ttte://frrumm1ddmfEYOoPPP
Qiyuan Liu, Qi Zhou, Rui Yang, and Jie Wang. Robust representation learning by clustering with bisimulation metrics for visual reinforcement learning with distractions. In Thirty-Seventh AAAI Conference Artificial Intelligence, pp. 8843-8851. AAAI Press, 2023b.
Martin Bertran Lopez, Walter Talbott, Nitish Srivastava, and Joshua M. Susskind. Efficient embedding of semantic similarity in control policies via entangled bisimulation. CoRR, abs/2201. 2022.
Chen Ma, Dylan R. Ashley, Junfeng Wen, and Yoshua Bengio. Universal successor features for transfer reinforeement learning. CoRR, abs/2001.0400 1.04255 2020.
Hiroshi Makino. Arithmetic value representation for hierarchical behavior composition. Nature Neuroscience, 26(1):140-149 2023.
Bogdan Mazoure, Ahmed M. Ahmed, R. Devon Hjelm, Andrey Koloboy, and Patrick MacAlpine. Cross-trajectory representation learning for zero-shot generalization in RL. In The Tenth Interna- tional Conference on Learning Representations. OpenReview.net, 2022.
Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Andrei A. Rusu, Joel Veness, Mare G. Belle- mare, Alex Graves, Martin A. Riedmiller, Andreas Fidjeland, Georg Ostrovski, Stig Petersen, Charles Beattic, Amir Sadik, Ioannis Antonoglou, Helen King, Dharshan Kumaran, Daan Wierstra, Shane Legg, and Demis Hassabis. Human-level control through deep reinforcement learning. Nat., 518(7540):529-532 2015.
Yingjun Pei and Xinwen Hou. Learning representations in reinforcement learning: An information bottleneck approach. CoRR, abs/191 105695, 2019.
14

---

Published as a conference paper at ICLR 2025
Han Qi. Yi Su, Aviral Kumar, and Sergey Levine. Data-driven offline decision via invariant representation learning. In Advances in Neural Information Processing Systems 35, 2022.
Roberta Raileanu and Rob Fergus. Decoupling value and policy for generalization in reinforcement learning. In Marina Meila and Tong Zhang (eds.), Proceedings of the 38th International Conference Machine Learning, volume 139 of Proceedings of Machine Learning Research, pp. 8787-8798. PMLR, 2021.
Roberta Raileanu, Maxwell Goldstein, Denis Yarats, Ilya Kostrikov, and Rob Fergus. Automatic data augmentation for generalization in reinforcement learning. In Mare Aurelio Ranzato, Alina Beygelzimer, Yann N. Dauphin, Perey Liang, and Jennifer Wortman Vaughan (eds.), Advances in Neural Information Processing Systems 34, pp. 5402-5415, 2021.
Stuart Russell and Andrew Zimdars. Q-decomposition for reinforcement learning agents. In Machine Learning, Proceedings of the Twentieth International Conference, pp. 656-663. AAAI Press, 2003,
Andrew M. Saxe, Yamini Bansal, Joel Dapello, Madhu Advani, Artemy Kolchinsky, Brendan D. Tracey, and David D. Cox. On the information bottleneck theory of deep learning. In 6th International Conference on Learning Representations. OpenReview.net, 2018.
John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. CoRR, abs/1707 .63447 2017.
Zhihao Shi, Xize Liang, and Jie Wang. LMC: Fast training of GNNs via subgraph sampling with provable convergence, In The Eleventh International Conference on Learning Representations, 2023. URLLttts:://ppnteee
n forum? d= 5VBBA 91N6n
Zhihao Shi, Jie Wang, Fanghua Lu, Hanzhu Chen, Defu Lian, Zheng Wang, Jieping Ye, and Feng Wu. Label deconvolution for node representation learning on large-scale attributed graphs against learning bias. IEEE Trans actions on Pattern Analysis and Machine Intelligence, 46(12):11273- 11286, 2024. doi: (0.1109/TPMMI.202 20t...244..4444000
Zhihao Shi, Jie Wang, Zhiwei Zhuang, Xize Liang, Bin Li, and Feng Wu. Accurate and scalable graph neural networks via message invariance. In The Thirteenth International Conference Learning Representations, 2025. URL
net/ wtt.:/orrmmreeeee wtwpeeewwneeee/orrmmmddeeeeeeeee
hcmF P
Xingyou Song, Yiding Jiang, Stephen Tu, Yilun Du, and Behnam Neyshabur. Observational overfitting in reinforcement learning. 8th International Conference on Learning Representations, 2020.
Austin Stone, Oscar Ramirez, Kurt Konolige, and Rico Jonschkowski. The distracting control suite- A challenging benchmark for reinforcement learning from pixels. CoRR, abs/2101.02777 2021.
Richard S Sutton and Andrew G Barto. Reinforcement learning: An introduction. MIT press, 2018.
Richard S. Sutton, Doina Precup, and Satinder Singh. Between mdps and semi-mdps: framework for temporal abstraction in reinforcement learning. Artif. Intell., 112(1-2):8812111 1999.
Hongyao Tang, Zhaopeng Meng, Guangyong Chen, Pengfei Chen, Chen Chen, Yaodong Yang, Luo Zhang, Wulong Liu, and Jianye Hao. Foresee then evaluate: Decomposing value estimation with latent future prediction. In Thirty-Fifth AAAI Conference on Artificial Intelligence, pp. 9834-9842.
AAAI Press, 2021.
Yuval Tassa, Yotam Doron, Alistair Muldal, Tom Erez, Yazhe Li, Diego de Las Casas, David Budden, Abbas Abdolmaleki, Josh Merel, Andrew Lefrancq. Timothy P. Lillicrap, and Martin A. Riedmiller. Deepmind control suite. CoRR, abs//101. 01111180100090, 2018.
Naftali Tishby and Noga Zaslavsky. Deep learning and the information bottleneck principle. In 2015 IEEE Information Theory Workshop, pp. 1-5. IEEE, 2015.
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, pp. 0026-5033. IEEE, 2012.
15

---

Published as a conference paper at ICLR 2025
Aäron van den Oord, Yazhe Li, and Oriol Vinyals. Representation learning with contrastive predictive coding. CoRR, 0tbs/1807.00 $97700374, 2018.
Harm van Seijen, Mehdi Fatemi, Romain Laroche, Joshua Romoff, Tavian Barnes, and Jeffrey Tsang. Hybrid reward architecture for reinforcement learning. In Advances in Neural Information Processing Systems 30, pp. 5392-5402, 2017.
Jie Wang, Rui Yang. Zijie Geng, Zhihao Shi, Mingxuan Ye, Qi Zhou, Shuiwang Ji, Bin Li, Yongdong Zhang, and Feng Wu. Generalization in visual reinforcement learning with the reward sequence distribution. CoRR, abs/2302.09601. 2023a.
Jie Wang. Zhihai Wang, Xijun Li, Yufei Kuang, Zhihao Shi, Fangzhou Zhu, Mingxuan Yuan, Jia Zeng, Yongdong Zhang, and Feng Wu. Learning to cut via hierarchical sequence/set model for efficient mixed-integer programming. IEEE Transactions on Pattern Analysis and Machine Intelligence, pp.
1-17, 2024a. doi: 10. 10.110992MM2222233337777
Jie Wang, Mingxuan Ye, Yufei Kuang, Rui Yang, Wengang Zhou, Houqiang Li, and Feng Wu. Long-term feature extraction via frequency prediction for efficient reinforcement learning. IEEE Transactions Puttern Analysis and Machine Intelligence, 2025a.
Kaixin Wang, Bingyi Kang, Jie Shao, and Jiashi Feng, Improving generalization in reinforcement learning with mixture regularization. In Hugo Larochelle, Marc' Aurelio Ranzato, Raia Hadsell, Mlrii-Florria Balcan, and Hsuan-Tiee Lin (eds.), Advances in Neural Information Processing Systems 33, 2020.
Shuo Wang, Zhihao Wu, Jinwen Wang, Xiaobo Hu, Youfang Lin, and Kai Lv. How to learn domain- invariant representations for visual reinforcement learning: An In Proceedings of the Thirty-Third International Joint Conference on Artificial Intelligence, pp. 1389-1397. ijcai.org, 2024b.
perspective.
Zhihai Wang, Jie Wang, Qi Zhou, Bin Li, and Houqiang Li. Sample-efficient reinforcement learn- ing via conservative model-based actor-critic. In Thirty-Sixth AAAI Conference on Artificial Intelligence, pp. 8612-8620. AAAI Press, 2022.
Zhihai Wang, Xijun Li. Jie Wang, Yufei Kuang, Mingxuan Yuan, Jia Zeng, Yongdong Zhang, and Feng Wu. Learning cut selection for mixed-intege linear programming via hierarchical sequence model. In The Eleventh International Conference on Learning Representations, 2023b.
Zhihai Wang, Taoxing Pan, Qi Zhou, and Jie Wang. Efficient exploration in resource-restriited reinforcement learning. Proceedings of the AAAI Conference on Artificial Intelligence, 37(8): 10279-10287, Jun. 2023c.
Zhihai Wang, Jie Wang, Dongsheng Zuo, Yunjie Ji, Xinli Xia, Yuzhe Ma, Jianye Hao, Mingxuan Yuan, Yongdong Zhang, and Feng Wu. A hierarchical adaptive multi-task reinforcement learning framework for multiplier circuit design. In Forty-first International Conference on Machine Learning. PMLR, 2024c.
Zhihai Wang, Jie Wang, Xilin Xia, Dongsheng Zuo, Lei Chen, Yuzhe Ma, Jianye Hao, Mingxuan Yuan, and Feng Wu. Computing circuits optimization via model-based circuit genetic evolution. In The Thirteenth International Conference on Learning Representations, 2025b.
Guofei Xiang, Songyi Dian, Shaofeng Du, and Zhonghui Lv. Variational information bottleneck regularized deep reinforcement learning for efficient robotic skill adaptation. Sensors, 23(2):762, 2023.
Eliot Xing, Abhinav Gupta, Sam Powers, and Victoria Dean. Kitchenshift: Evaluating zero-shot generalization of imitation-based policy learning under domain shifts. In NeurlPS 202/ Workshop on Distribution Shifts: Connecting Methods and Applications, 2021.
Rui Yang, Jie Wang, Zijie Geng, Mingxuan Ye, Shuiwang Ji, Bin Li, and Feng Wu. Learning task-relevant representations for generalization via characteristic functions of reward sequence distributions. In The 28th ACM SiGKDD Conference on Knowledge Discovery and Data Mining, pp. 2242-222222222222222222 ACM, 2022.
16

---

Published as a conference paper at ICLR 2025
Denis Yarats, Ilya Kostrikov, and Rob Fergus. Image augmentation is all you need: Regularizing deep reinforcement learning from pixels. In 9th International Conference on Learning Representations. OpenReview.net, 2021.
Denis Yarats, Rob Fergus, Alessandro Lazaric, and Lerrel Pinto. Mastering visual continuous control: Improved bata-augmented reinforcement learning. In The Tenth International Conference on Learning Representations. OpenReview.net, 2022.
Mingxuan Ye, Yufei Kuang, Jie Wang, Rui Yang, Wengang Zhou, Houqiang Li, and Feng Wu. State sequences prediction via fourier transform for representation learning. In Alice Oh, Tristan Naumann, Amir Globerson, Kate Saenko, Moritz Hardt, and Sergey Levine (eds.), Advances in Neural Information Processing Systems 36, 2023.
Zhecheng Yuan, Zhengrong Xue, Bo Yuan, Xueqian Wang, Yi Wu, Yang Gao, and Huazhe Xu. Pre-trained image encoder for generalizable visual reinforcement learning. In Sanmi Koyejo, S. Mohamed, A. Agarwal, Danielle Belgrave, K. Cho, and A. Oh (eds.), Advances in Neural Information Processing Systems 35, 2022.
Zhecheng Yuan, Sizhe Yang, Pu Hua, Can Chang, Kaizhe Hu, Xiaolong Wang, and Huazhe Xu. RI-vigen: A reinforcement learning benchmark for visual generalization. CoRR, abs/2307.10224, 2023.
Amy Zhang, Clare Lyle, Shagun Sodhani, Angelos Filos, Marta Kwiatkowska, Joelle Pineau, Yarin Gal, and Doina Precup. Invariant causal prediction for block mdps. In Proceedings of the 37th International Conference on Machine Learning, volume 119 of Proceedings of Machine Learning Research, pp. 11214-11224 PMLR, 2020.
Amy Zhang, Rowan Thomas McAllister, Roberto Calandra, Yarin Gal, and Sergey Levine. Learning invariant representations for reinforcement learning without reconstruction. In 9th International Conference on Learning Representations. OpenReview.net, 2021.
Ruijie Zheng, Xiyao Wang, Yanchao Sun, Shuang Ma, Jieyu Zhao, Huazhe Xu, Hal Daumé III, and Furong Huang. Temporal latent action-driven contrastive loss for visual reinforcement learning. Advances in Neural Information Processing Systems, 36, 2024.
Qi Zhou, Jie Wang, Qiyuan Liu, Yufei Kuang, Wengang Zhou, and Houqiang Li. Learning robust representation for reinforcement learning with distractions by reward sequence prediction. In Robin Evans and Ilya Shpitser (eds.), Uncertainty in Artificial Intelligence, volume 216 of Proceedings of Machine Learning Research, pp. 2551-2562. PMLR, 2023a.
Weitao Zhou, Zhong Cao, Yunkang Xu, Nanshan Deng, Xiaoyu Liu, Kun Jiang, and Diange Yang. Long-tail prediction uncertainty aware trajectory planning for self-driving vehicles. In 2022 IEEE 25th International Conference on Intelligent Transportation Systems (ITSC), pp. 1275-1282. IEBE, 2022.
Weitao Zhou, Zhong Cao, Nanshan Deng, Kun Jiang, and Diange Yang. Identify, estimate and bound the uncertainty of reinforcement learning for autonomous driving. IEEE Transactions on Intelligent Transportation Systems, 24(8):7932-7942, 2023b.
Zhuangdi 7hu, Kaixiang Lin, Anil K. Jain, and Jiayu Zhou. Transfer learning in deep reinforcement learning: A survey. IEEE Trans. Pottern Anal. Mach. Intell, 45(11):1334441333 13362, 2023.
17

---

Published as a conference paper at ICLR 2025
APPENDIX
A RELATED WORK OF TASK DECOMPOSITION
Value Decomposition. Hierarchical RL aims to solve complex tasks into layered sub-goals, enabling agents to tackle challenges at varying levels of abstraction for efficient and scalable learning (Wang
et al. 2023b 2024a; Ling et al., 2024 Kuang al. 2024 Liu et al. 2024 2025 2023a). Some methods suggest that reusing learned skills from simple tasks can effectively improve the sample efficiency in solving the complex composite tasks (Sutton al. 1999. Dietterich 2000 Russell & /Zimdars 2003 Makino 2023). Thus, for a complex sequential decision-making task that can be divided into series of simple individual sub-tasks, these approaches estimate the action-value
function for each sub-task to learn these reusable skills.
Inspired by this hierarchical RL, some methods propose the reward decomposition for composite tasks, which decomposes the one-step rewards into a series of sub-rewards instead of breaking
down the composite task (van Seijen et al., 2017: Lin al.) 2019 Grimm Singh 2019; Fatemi & Tavakoli 2022). Specifically, under the assumption that one-step rewards of composite tasks are linear combinations of independent sub-rewards, these methods decompose each one-step reward either manually or automatically (i.e., leveraging deep learning techniques). Then, they combine the action values for each sub-task associated with each sub-reward additively to estimate the overall action-value function, thus efficiently addressing the complex composite task.
In this article, although the reward representations extracted by ROUSER are similar to the sub- rewards from reward decomposition methods, they do not rely the prior assumption that one-step rewards are linear combinations of independent sub-rewards.
Successor Features. Successor features have garnered attention in the field of transfer RL, which allows flexible knowledge transfer across tasks with different reward functions. Barreto et al. (2017) assume set of tasks with different reward functions but similar environment dynamics, and thus propose the successor features, which represent the features from environment dynamics involving the long-term sequential information. To learn such features, [Barreto al. (2017) 2018 2020);
Ma et al. 20200; Carvalho et al. (2023) first decouple rewards into the component associated with environment dynamics and the weight vector associated with different tasks. Then, they can use the reward component to learn successor features from the environment dynamics. By using the successor features and the weight vector, they can achieve the transfer learning. Moreover, based on successor features, Abdolshah et al.) (2021) propose transferring knowledge across different environments by using a Gaussian distribution to capture the changes environment dynamics.
Our study differs from these methods, which do not take into account robustness against visual distractions. Also, the successor features learned by the aforementioned methods often involve irrelevant information from environment dynamics in the VRL generalization challenge. Instead, we use one-step rewards to extract long-term information from action values and leverage the IB framework to discard irrelevant features. It is worth NOTING that in all the aforementioned successor feature works, fine-tuning the successor features is essential when encountering unseen testing environments. In contrast, our approach focuses zero-shot generalization without any fine-tuning.
B IMPACT STATEMENTS AND LIMITATIONS
This paper proposes a novel approach called ROUSER to advance the field of generalization in VRL, enhancing the potential of agents in real-world applications. Although our primary focus is on technical innovation, we recognize the potential societal consequences of our work, as generalization in VRL has the potential to influence various domains such autonomous driving and the financial field. We are committed to ethical research practices and attach great importance to the social implications and ethical considerations in the development of generalization research in VRL.
is essential for our approach to use the artificially designed rewards for robust action-value representation learning. We look forward advancing the representation learning in VRL using inverse RL techniques that can derive the reward function from data. This will enable us to introduce framework that supports autonomic iterative optimization. Moreover, our work may inspire some
18

---

Published as a conference paper at ICLR 2025
studies about the message invariance for message passing based graph neural networks (Shi et al.
2023 2025: 20044.
C
PROOFS
C.I
PROOF OF THE MAPPING
In this subsection, we prove that the action-value representation is linearly correlated with the infinite sequence of reward representations.
Proof. Based on Equations
3 and the definition of the action-value function, we have
Q (or, ar) Ex,p
(17)
Ex
7(',, h(octt)
(18)
w' En.
Σ
7sh hhlotee asti)
(19)
(w', HT (ot, Therefore, derive action-valus representation H" (", (o,,a.).
(20)
C.2 PROOF OF THE CONVERGENCE OF ACTION-VALUE REPRESENTATION LOSS
In this section, we provide a detailed derivation of the convergence action-value representation. ROUSER leverages the reward representations to introduce auxiliary objective function with recursive form. Similar to the TD-learning of action-value functions, we prove that Equations and 5 can be reformulated a contraction mapping T. Here, define Z" generic function 2
x A + 2, and we define the mapping T Z Z.
Theorem There exists contraction mapping T Z 2 defined
Yor,)or,o at) = =x(0r, ax) + Er., 1Z (0 0c+1) at]
Proof. Let (o,, (o)) denote the L1 value of the d di-mmmensoonn of 2 (ot, az), where < d S D. Due to the properties of contraction mappings (Sutton & Barto 2018 Zhou al., 2023b), we can apply the operator T to compute the target auxiliary loss function until convergence in tabular settings. operator a contraction in the sup-norm, which we provide as follows.
TT
TZ5 1oo 7 ||2
2 1lo0
(21)
We provide detailed proof of Equation 21 |
TZZ loc max 12 + E.,,p 7 (0++1,
- E..p [ (ot+1, a441)oo adll
Y max [Bm,p [ (o++1, (14144110l00 arl E. [2, (ot+1, a++1) (141))ooo lo adlaa
7 max [E.r, [ (0++1, 0+41 23 (0++1, ax+ (0t, azlla
S 7 max Fx, 1 (0++1, 4441) Z (ot+1, a1 at]
$ 7 max
7 (0 4++1 27 (ot+1. a+11)
00
a111EA
7 || Z xx.
Based on Equation[1] 21] we can derive the convergence of auxiliary loss similar to action-value function estimation. So apply T to compute the auxiliary loss until convergence in tabular settings.
19

---

Published as a conference paper at ICLR 2025
C.3 INFORMATION BOTTLENECK
Here, we provide derivation of the formulas involved in the IB techniques in our ROUSER approach. We first list the mathematical formulas used in our proof.
(P.1) KL Divergence
p(x)
Dice (p(x)]la(x)) Ex
log
(22)
g(xx)
(P.2) Mutual Information
log y)x, p(x)pply)
I(x;y) = Ex.y
= Dxc (pxx
(23)
(P.3) Chain rule of Mutual Information
I(x,y:z) I(y:z) + zly).
(24)
(P.4) Symmetry of Mutual Information
I(x;y) = x(y;))
(25)
Then, for Equation 6 leverage (P.2) expand the computation of mutual information and conditional mutual information as follows.
p(r, z) p(r)p(z)
I(r;
Er,x
log
p(z)p(r(z)
Er,
log
(26)
p(r)p(z)
log p(r|z)
Er
p(r)
log p(u, z/r) p(u[r)p(z(r)
T(u:z/r) = Eu,zr
log p(ulr)p(z(u.rr)
Eu,zr
p(u|r)p(z(r)
(27)
log p(z u, r p(z En,x.
p(z u)
Eu,zr
log
p(z
Here in since z is the representation of we have p(z) u, p(z/u). As discussed in Equation g we can introduce a new feature to provide an upper bound to estimate the conditional mutual information. The details of the proof in Equation 8 are follows.
I(u; z/r)
p(z (u)
Eu.rr
log
p( [
log p(z (u) p(k(r)
=Euz.
p(k[r) p(z|r)
p(z lu)
log p(z(r)
Eu.zr
log p(k[r)
Eu.zr
p(k[r)
(28)
p(z u))
log p(z(r)
Erjux
Eu,z log
1E z
Ez,r
p(k[r)
p(k(r)
log p(z|u) p(k(r)
log p(z(r)
Eux
Ez,r
p(k(r)
DrC (p(z/u)llp(k)r)) Drc (((z)r)lp((((r)
< Dicc (p(z(u)ll((kir) = Dxcc (NalllN.)
20

---

Published as a conference paper at ICLR 2025
Since the features z ~ Ny and k ~ Ny, we can directly use the Gaussian distribution /V to replace the probability density function Note that the equality holds the two distributions coincide. Similar to [(u; z/r), [rr k/u) is upper bounded by Dxc (N.)
We can also leverage k to provide lower bound for the mutual information:
I(r;2)
(13)
r(k, I(z;k)r)
(29)
= r;)
(30)
P3) I(k;2) + r(z;r(k)
(31)
2 Z(k; z).
(32)
Since is the representation of r, we have I(z; k = 0, and we thus derive the equality with of
Equation 30 Moreover, the equality in Equation 32 holds if [(z; r/k) = 0, i.e. the representation
captures all the information associated with the reward r.
Based on Equation
we can further derive the whole optimization objective Lebb as:
[(u;2/r) + I(r; k|w)] - BI(r;2) Lubb
Dxc (NellNelNn) + Dxc (N,l 3I(k;2)
(33)
Dsxcc(NelWG) - -II(k;2)
By reparameterizing the objective, derive the final loss function:
Linto (v,, 1)) Es -I(z;k) + DDcc ( (NollNn))
(34)
C.4 PROOF OF THE ACTION VALUE BOUND
We provide the detailed proof of the Theorem 4.4
Theorem 4.4 Let Z be learned robust action-value representation from any observation-action
pair (o, a)
OxO Q O x A + R be the true action-value function of a policy æ in the
environment E l 2 + R be the optimal linear mapping on the representation spuce, and be a bound of estimation error for each compressed reward representation z, i.e., (f2(2) - r For
any (o, a) E O x A and E E, we have
1 I-7
< (f(
Proof.
(0, - fz (Z(o
E
E
(2)
10c
(n. f2)))
8
Re
7' ||r
[2(=)]1
R
1
21

---

Published as a conference paper at ICLR 2025
D ARCHITECTURE AND PSEUDOCODE
As we directly use the embeddings from critic's center layer as robust action-value representations, we do not need to change the architecture of commonly used VRL algorithms. Specifically, we follow the traditional VRL actor-critic architecture use 3-layer ferw-rorward ConvNet with no residual connection as the encoder. Then, we apply a 3-layer MI.P with hidden size 1024 as each critic, and we use Equation| 15 to regulate the embeddings of the second layer of each critic. The actor uses the same architecture the critic. Moreover, use additional 2-layer MLP with hidden size 1024 as the reward model to to output the compressed reward representations 2, as well as yj to output the representations k from one-step rewards. As the inputs image observations, the reward model eb after the shared encoder for the actor and critic. We provide the architecture and the algorithm of ROUSER in Figure and Algorithm respectively.
Algorithm 1 ROUSER in a general actor-critic framework
Initialize the critic network Qe with parameters actor network The with parameters Φ, reward model v;, and the linear mapping f-
Initialize target network Qa with weight 0 <- 8. Initialize the replay buffer B. for t 1,..., do if f < To then
Randomly select under o. Execute to obtain and o'.
else
a (0
o' d'oplo'oo,a)
Sample batch to, a, r, or) from B.
Update:
Xy V foood
linear network reward network critic network actor network target network
v 1/ AN Cpred Linto)
f x- 8 Ao Vo (Crobust Ceruic)
@ + Φ As V &aaace
1 + r + (1 T)
end if
+ BU fo
end for
E EXPERIMENTS
E EXPERIMENTS ON DMC
Environment Setting. Color Distractions. On DMC benchmark, we apply the similar treatment of dynamic color distraction (Stone al. 202I) the objects of the environments. Specifically, uniformly sample the color z0 ~ U(x - B, x is the origin color in DCS, and 3 is hyperparameter. We leverage dynamic setting where the color 2; can change to rt+1 clip(2 (2, 24 3, x; + 3), where T Nxx,, 0.03 3). We train
3) for each channel at the start of each episode, where
the agents on the environment with weak dynamic color distractions (3 0.2). Then, we evaluate all agents in the test environments with strong dynamic color distractions (3 0.5).
Background Distractions. We adopt the dynamic background settings from Stone al. (021). To establish different training environments, we utilize N videos from the DAVIS 2017 training set, where represents the number of training environments. Each environment uses one video as the background and randomly samples scene and frame from the video at the start of every episode. Additionally, we set Pog = 1.0, indicating that we use the distracting background instead of the original skybox. For evaluation, apply 30 videos from the DAVIS 2017 validation dataset the unseen backgrounds. In each episode of the test environment, we randomly select one of the 30 dynamic backgrounds.
22

---

Published as a conference paper at ICLR 2025
Figure 6: The dynamics color and background distractions using in Section 5.1 The first row illustrates the dynamic color changes, and the second row shows one of dynamic backgrounds.
Cumulative Reward
1c2
ROUSER-DAQU2
ROUSER-SRM
5
DrQU2
w
SRM
Loss
CRESP
2
MSE
weighs 0.01 wright weighs-n.5
33
weight: 1
2
0.00
0.25
0.50
0.75
1.00
2
4
Environment Steps
l66
lel
Figure 7: On cheetah run task under color changes, the left of this figure shows the curves of MSE loss for capturing long-term information, and the right illustrates the performance of weight
Network. For the shared encoder, we use the default setting of baselines combined with our approach. For example, in ROUSER-SRM, we use shared pixel encoder with 11 convolutional layers to extract the image information, and employ 3 x kernels and 32 filters with stride of 2 for the first convolutional layer and 1 for others. After each convolution, following the commonly used setting of VRL methods (e.g., DrQ, DrQv2, and SRM), we apply a fully connected layer to output 50-dimensional representations normalized by LayerNorm and tanh activation. Note that the actor and critic each have fully connected layer.
After the fully connected layer, the actor network is parameterized with fully connected layers using ReLU activations up until the last layer. The critic network employs Double Q-learning technique, where each Q-value learned using 3-layer fully connected network, similar the one used in the actor network. The output dimension of these hidden layers in the actor and the critic network is 1024. Moreover, the gradients the shared pixel encoder are computed through the critic's optimizer rather than the actor's
The reward model also has fully connected layer to output 50-dimensional representations normal- ized by LayerNorm and tanh activation from the outputs of the shared encoder. After the fully connected layer, the reward model has 3-layer MLP. The output dimension these hidden layers in the reward model is also 1024. The final outputs of the reward model is twice the dimension of compressed reward representations, half of which is the mean and half of which is the variance. Thus, use the outputs to model the Gaussian distribution of compressed reward representations. Then, we use a linear layer, mapping compressed reward representations into the corresponding one-step rewards. Moreover, use linear layer get the representations k from one-step rewards.
23

---

Published as a conference paper at ICLR 2025
1.0 1e1
color ball, cup catch
1.0
ls1
color cartpole swingup
10
color cheetah run
0.6
0.6
0.0
0.0
1.0
163
color finger spin
1.0
163
color walker walk
1.0
163
voLor reacher easy
0.5
0.6
0.6
0.4
0.4
0.0
0.0
0.0
1.00
163
cup,
1.00
163
video carpolo swingup
1c2
video chestah
0.75
0.75
0.50
0.50
0.25
0.25
0.00
0.00
0
le5
1.00
162
flnger spin
1,00
163
video walker walk
1,00
161
video rencher easy
0.75
0.75
0.75
0.50
0.50
0.50
0.25
0.25
0.25
0.00
0.00
0.00
165
ROUSER-SRM
CRESP
ROUSER-DIQU2
CURL
SRM
RAP
DrQv2
TACO
Figure 8: Training curves over six seeds on DMC benchmark with color and video buckground distractions. Each point from the curves is evaluated over ten episodes on test environments with unseen distractions.
Implementations. In the implementation of VRL, the multi-step reward is a commonly used trick
without any theoretical guarantee: r(Ot, a.) is (f((otii, 04+s), where f(:,.) is the reward function. Many VRL methods (Yarats et al., 2022] Huang et al. 2022 Hessel et al. 2018) use N 3. In our approach, we do not change the default setting when combined with baselines. It is worth noting that we change the mode of our reward model. Specifically, we map the outputs of reward model into the N raw rewards instead of the multi-step reward, thus learning the compressed reward representations. Unlike CRESP, which predicts reward sequence distributions to update the shared encoder for representation learning, our reward model does not update the shared encoder with this mapping, although it establishes mapping from the shared encoder's outputs to the multi-step rewards. This is because that we believe the multi-step rewards with length of N do not involve much long-term robust information for sequential decision decision-making making. We demonstrate this in the left of Figure The experiment setting for this figure is the same as in Figure 4 (a). We use 200K collected data to learn 2-lerrr MLP projector that maps the fixed representations from CRESP and CAUSER-DIQVE to the long-term sequential rewards with the length of 300. The curves demonstrate that ROUSER effectively capture more long-term term robust information than CRESP.
E
Compute Resource. We use NVIDIA GeForce RTX 3090 GPUs for six tasks on DMC benchmark under backgrounds or color changes. Trials of DrQv2, SRM, D and ROUSER-SRM
24

---

Published as a conference paper at ICLR 2025
Table 1: Means and standard emrors six DMC environments with unseen dynamic color changes. Each result is averaged over 100 episodes at 500K environment steps using six random seeds.
ALGO
ball call_in_cup_cam catch
carppllesswnngup
cheetah_run
linger_spin
walker_walk
reacher_eee _
ROUSER-SRM
942.6 12.3
683.5 : 37.8
370.0 25.1
853,4 25,6
814,0 :: 19.2
754.8 39.7
ROUSER-DIQU2
780.9 = 60.3
461.4 15.6
371.6 23.7
424.2 = 49.8
417.7 = 22.7
154.6 = 52.8
SRM
826.0 + 125.6
640.2 15.8
299.5 1: 23.4
779.7 = 46.8
776.5 +: 27.8
659.4 4: 18.4
DrQu2
632.1 49.5
377.0 24.7
140.2 10.2
234.0 54.9
115.3 41.8
93.7 20.4
CRESP
57,8
410.6 15.2
87.
574.0
394.1
270.9 17.6
CURL
31.
46..6
279.2 +: 27.6
278.2 1 46.3
198.2 4: 37.4
TACO
360.5 1905:19222 192.8
385.9 26.3
138.7 17.2
130.8 45.7
177.8
= 318
122.9 = 28.2
RAP
149.2 = 34.2
304.4 44.4
90.8 = 57.6
1.5
86.5 + 62.5
209.8 59.7
Table 2: Means and standard errors on six DMC environments with unseen dynamic buckgrounds. Each result is averaged over 100 episodes at 500K environment steps using six random seeds.
ALGO
ball in catch
Surppll_swwnngpp
finger_spin
cheetah_run
walker_walk walk
reacher_ewse
ROUSER-SRM
868.9 47.0
822.4 + 10.8
776.8 + 24.0
372.9 + 10.8
894.0 + 12.6
901.5 + 18.2
ROUSER-DIQS
523.1 67.5
731.9 + 10.5
348.2 +: 174.1
408.0 + 22.1
358.6 +: 74.7
440.0 +: 70.9
SRM
818.1
807.7
527.5 + 27.2
22,9
823.1
889.2 14.2
DrQv2
474.3 = 72.6
717.8 15.0
309.7 = 113.8
359.2 = 19.2
237.7 + 106.9
404.5 :: 32.9
CRESP
665.3 = 75.5
6 20.0
778.1 62,9
327.4 + 20.1
794.1
667.7 +: 33.5
CURL
167.0 57.4
329.6 18.4
745.9
185.0 + 15.9
746.1 16.7
714.9 33.1
TACO
816.1 42.6
67.
737.2
59.5
505.4 151.2
198.8 512
RAP
76.1 + 14.6
170.8 21.7
6,9 = 4.0
96.7 12.4
115.1
150.9 + 16.4
on reacher_ean _aasy task from DMC under dynamic backgrounds are trained for 12.24. 10.83, 16.93, and 15.83 hours on average. Moreover, the agents of ROUSER-DrQv2 and ROUSER SRR require approximately 2800MB and 3600MB of memory using the batch size of 256, respectively.
Results. We compare our approach ROUSER against several SOTA method: (1) CURL (Laskin et al., 2020), which leverages contrastive learning to maximize the mutual information between
representations from observations and augmentations. (2) DrQv2 (Yarats al. [022]. which the prior state-of-the-art DRL algorithm for sample efficiency.
(3) SRM (Huang et al.) 2022).
which adopts augmentation in the frequency domain to facilitate the learning of robust policies.
RAP (Chen & Pan 2022). which effectively enhances the robustness of representations by leveraging the behavioral similarity. (5) CRESP (Yang et al.. 2022), which predicts the charactcristic function of reward sequences learn task-relevant representations for generalization. (6) TACO (Zheng et al. 2024), which learns state and action representations that encompass sufficient information for control improve sample efficiency. Moreover, we illustrate image observations from the environments with dynamic color and background distractions in Figure
We provide the means and standard errors 500K environment steps in Tables 1 and Note that Table [litts the detailed results of Figure under unseen dynamic color changes, and Table lists the detailed results of Figure Jnder unseen dynamic video backgrounds. As shown in Tables and 2. ROUSER outperforms other baselines in 11 out of 12 experiments. Even in finger_sppe spin task under dynamic backgrounds, our approach ROUSER is almost on par with the highest performance. In addition, we illustrate the training curves over six seeds in Figure g
Analysis of learned Representations for Captuing Long-Term Robust Information. In Section 5.3 we present results from task where learned representations are used to predict the average of future reward sequences over length of 300. These results indicate the effectiveness ROUSER in capturing long-term information.
Moreover, we conduct additional experiments with varying reward sequence lengths, including 50, 100, 300, and 500. We report the results in Table 3 which are averaged over three seeds on cheetah run task with unseen color distractions. These results, averaged over three seeds on cheetah run task with unseen color distractions. demonstrate that the representations learned by ROUSER are able to accurately predict reward sequences of different lengths, with performance improving as the reward sequences become longer.
25

---

Published as a conference paper at ICLR 2025
Reward Sequence
ROUSER-SRM
SRM
ROUSER-DIQV2
DrQv2
Length=50
0.0169 0.0051
0.0170 0.0056
0.0169 + 0.0053
0.0173 + 0.0029
Length=100
0.0213 + 0.0070
0.0210 + 0.0069
0.0205 + 0.0065
0.0211 + 0.0004
Length=300
0,0244 0.0122
0.0268 0.0179
0.0239 = 0.0182
0.0283 = 0.0029
Length=500
0.0294 = 0.0091
0.0383 + 0.0004
0.0287 = 0.0087
0.0319 + 0.0011
Table 3: MSE results and standard deviations with different reward sequence lengths over three seeds, including 50, 100, 300, and 500.
Game
PPO
MixReg
QR-DQN
HOUSER-QRDOR
bigfish
3.7 1.3
7.1 + 1.6
12.2 + 1.1
13.6 + 5.2
caveflyer
5.1 0.4
6.1 # 0.6
6.0
6.4 = 0.4
dodgeball
1.6 0.1
1.7 # 0.4
7.2 0.2
8.0 = 0.4
Table Average Procgen scores on test levels after training on 25M environment steps. The mean and standard deviation are computed using five results with different random seeds. We boldface the results that have highest means. This table corresponds the results Figure
E.2 EXPERIMENTS ON PROCGEN
Settings. Procgen benchmark consists of 16 procedurally generated games. Each of these games has procedurally generated levels which present agents with meaningful generalization chal- lenges (Raileanu al. 202T). All environments of these games use discrete 15-dimensional action space, and produce 64 x 64 x 3 RGB observations. Following Raileanu et al. (2021), we use the easy setting, where agents are learned on the training environments with 200 levels and tested on environments with unseen levels.
Implementations. Value-based VRL methods mainly focus on non-continuous control tasks. They often critic estimate the action-value function without actor/policy network. Thus, their critic inputs an image observation and outputs a vector of action values for all possible discrete actions. Note that the last layer of their critic maps the representations of image observations into the action values of all possible discrete actions. Therefore, in the combination of ROUSER and QR-DQN, for the action value of each possible discrete action, we replace its last layer mapping into 2-layer ensemble MLP, where the ensemble size is the number of discrete actions. The inputs of this 2-layer ensemble MLP are robust action-val value representations, and this MLP outputs the quantile values of the action-value function for each discrete action.
Moreover, we use the same design for the reward model s with 3-layer MLP. Specifically, the reward model uses the image observations inputs and outputs the compressed reward representations for each discrete action. uses actions from batch to select compressed representations for updating.
Results. We compare our approach with: (1) PPO (Schulman al. 2017 popular policy gradient baseline upon which many competitive methods are developed; (2) MixReg (Wang et al., 2020). regularization method applicable to both policy gradient and palue-based RI. algorithms to enhance generalization on Procgen; (3) QR-DQN, a prior SOTA based method on Procgen.
Table 4 lists the Procgen scores on test levels after training 25M environment steps using the batch size of 256. The mean and standard deviation are averaged over five runs with different random seeds. We adapt the tables from Raileanu & Fergus (02I). These results demonstrate that ROUSER can be also applicable to value-based RL algorithms, outperforming several previous methods.
E.3 EXPERIMENTS ON MuJoCo
To demonstrate that ROUSER can not only improve robustness VRL but also enhance sample efficiency in traditional state-based RL (i.c., RL using vector state), further conduct experiments on MuloCo (Todorov et al., 2012) using vector states. We combine our approach with DDPG (Lillicrap al. 2016) and SAC (Haarnoja et al. 2018), and further compare against the
26

---

Published as a conference paper at ICLR 2025
Cumulative Reward
14
163
HalfCheetah
6
163
Walker
40
1e2
Hopper
12
100
30
8
20
6
2
0
4
6
10
2
4
6
10
2
6
10
Cumulative Reward
L63
Humanoid
Ant
JnveriedDDuubbePeedulllu
10
6
3.
4
-2
2
6
8
1.0
6
8
1.
4
6
10
Environment Steps
le5
Environment Steps
163
Environment Steps
le5
ROUSER-SAC
VDFP
ROUSER-DDED
SAC
DDPG
Figure 9: Training curves during the training process with 1M steps six random seeds.
Table Hyperparameters of ROUSER on DMC and Procgen benchmarks.
Parameter
DMC Value
Procgen Value
Optimizer
Adam
Adam
learning rate of critic/action-value network
1 10-1
2.5 10-4
learning rate of robust action-valte value representations
1. 10-5
2.5 10-5
learning rate of actor network
1 10
None
learning rate of reward model target smoothing coefficient update frequency of target network dimension of robust action-value representations D
1 1.10-
2.5 10-4
0.01
0.005
2
1
512
512
batch size total environment steps T a scheduler start value scheduler end value reward model update times for each step nstep
256
512
106
2.5 10?
10-3
10-2 2
10-2
10-1
5
5
3
3
SOTA task decomposition methods, VDFP (Tang et al. 2021). We illustrate training curves (mean and standard deviation) on six tasks from MoJoCo during the training process in Figure 9 Each result is run for million time steps over six random seeds. Figure shows that ROUSER outperforms baselines, achieving an average improvement of +18.7%.
E.4 EXPERIMENTS ON DMC-GB
Environment Setting. To comprehensively evaluate our approach, we have conducted additional experiments on six DMC-GB (Hansen Wang 202T) tasks across all their settings (i.e., color_easy, color hard, video _easy, and video_hard) hard). All these experiment settings are the same as Hansen & Wang (2021) and Hansen et al. (2021) Moreover, we compare our approach against several SOTA generalization VRL methods, including PIE-G Yuan al. 2022). SVEA (Hansen et al. 202 l). DrG (Ha et al. 2023), SECANT (Fan et al., 2021). SGQN (Bertoin et al.) 2022). and MIIR (Wang et al. 2024b). Most of the baseline results report are taken directly from their respective papers. For cases where results are not provided, we ran experiments using the available source code and hyperparameters under three random seeds.
27

---

Published as a conference paper at ICLR 2025
Table 6: Hyperparameters of ROUSER MuloCo benchmark for sample efficiency.
Parameter
Value
Optimizer
Adam
learning rate of critic network
1 10
learning rate of robust action-value representations
1 10-1
learning rate of actor network
1 10-4
learning rate of reward model target smoothing coefficient update frequency of target network
5 10
0.01
2
dimension of robust action-value representations D batch size
512
256
total environment steps scheduler start value
106
10-3
scheduler end value nstep
10
Table 7: Results and standard deviations DMC-GB under the color easy and color hard settings. Results of our approach are averaged over three random seeds.
color_easy
color_easy
KOUSER-SRM
PIE-G
DrG
SVEA 959 +
SGQN
TIA
SECANT
ball ball_in_enn cap_catch
963
955 +
831 + 92
907 + 71
652 + 402
perpoll_swinngue
841 + 20
624 + 52
701 + 43
826 + 20
598 + 92
483 + 62
cheetah_run
616 13
429 12 12
375 1 31
587 1.
304 1 43
281
finger_spin
933
845
+ 79
892
628 + 46
410 49
walker_wald
40
33
907
845 +
780
walker_stand
968
20
963
895
color_hard _ard
ROUSER-SRO
PIE-G
DrG
SVEA
SGQN
TIA
SECANT
ball iin_cuu__atcchhhhhpppppippaaaeai
964 =
960 1
607 1: 46
961 7
905 1
575
375
958 1
cartpole_swinget
799
520 + 69
523 +
837
540 7-
407
173
866 + 15
cheetah run
562
376 27
219
456
277
249
582 - 64
finger_spin
979
10
758 124
461
408
115
walker_wald
916
824 92
725
134
760 145
692 153
616
856 1
walker_stand
962
948 + 15
731
+: 21
942 +: 26
905 +:
841
939 7
voly 1. We rum PIE-G, DrG, and SGQN on all DMC-GB tasks across color easy and color hard settings.
2. We run DrG on finger_spin and walker_stane tasks across video_easy and video video_hard hard settings.
easy setting and cheetah run task in 3. We run SVEA all DMC-GB tasks in the color_easy the color hard setting.
4. We run TIA on all DMC-GB tasks in all four settings.
Note that since SECANT and MIIR did not release their source code, we only include the results reported in their respective papers.
Results. We provide all results in Tables and These results indicate that ROUSER outperforms the aforementioned baselines in 17 out of 24 settings, demonstrating ROUSER'S 's robustness and effectiveness in DMC-GB tasks.
E.5 HYPERPARAMETERS
Hyperparameters. We provide the hyperparameters for DMC, Procgen, and MuJoCo in Tables 6 In DMC benchmark, the hyperparameter in the loss function Linto can be tuned from a small value 3 3 to 10-2 by an exponential scheduler.
We the reproduced version of TIA provided http://gthhua.comm//cmmnnngnn
28

---

Published as a conference paper at ICLR 2025
Table 8: Results and standard deviations DMC-GB under the video_easy and video hard settings.
Results of our approach are averaged over three random seeds.
video_ easy
ROUSER
PIE-G
DrG
SVEA
SGQN
MIlR
TIA
SECANT
ball ctp_ii__uppppeee ca__cup_cttlll catch
979 +: 21
922 2:
701 : 36
871 +: 22
950 + 24
973 =
664 416
903 + 49
cartpole swingup
830 +:
587 61
761 28
=
326 139
38
chectah_ run
397 30
287 20
547 21
+ 20
308 34
393
225 = 41
428 = 70
finger_spin
970 6
837 + 107
751 + 4:43
808 + 3.33
956 + 26
978 +
295 + 16
861
102
walker_walk
925 + 17
871 + 22
902 + 23
819 71
9
919 +: 30
696 245
842 + 47
walker_stamd
973 1 3
957 + 12
910
961 8
955 1
971 :
903 1: 18
932 1 15
video hard
ROUSER
PIE-G
DrG
SVEA
SGQN
MIR
TIA
SECANT
ball Sall_in_cup__co catch
940 1
786 4.
635 1: 26
403 1:3::: 174
782 5.
929 1
314 1: 36
curtpolesswingup
401
cheetah
270 + 12
17
3:
135
-
=
finger
937 + 10
79
437 = 61
335 - 58
822 24
956 = 18
182 ::55
walker_ walk
868 + 59
600 100+:22
782 + 37
377 + 93
739 21
821 + 58
166 + 60
walker_stand stand
968 4
852
819
+ 46
851
24
965 +
673 142
Hyperparameter Search. In the right of Figure 7 we illustrate the results (mean and standard deviation) of ROUSER-DrQV2 on cheetah run task from DMC benchmark under color changes for searching the learning rate of robust action-value representations. Based the learning rate i = 0.0001, we introduce the weight such that I e Each curve is averaged over three seeds.
Moreover, in Figure of the main text, we report the results without standard errors of different a with three random seeds. Here, provide the detailed results (refer to Table 9)
Table 9: The mean and standard error of the results in Figure
scale-0.3
scale-0.4
scale-0.5
scale-0.6
0.01-0.1
649.7 + 9,1
488.0 + 15.7
363.0 + 17.3
316.1 # 9.1
0.001-0.01
663.7 + 7.5
524.4 14.5
398.4 19.4
283.3 # 11.2
0.0001-0.001
683.1 # 2.2
510.3 16.8
367.9 12.5
251.4 # 11.5
29


Reviews:
[{"review_id": "KivcZPBvRL", "reviewer": "Reviewer_M8ko", "summary": "This paper proposes a novel approach termed Robust Action-Value Representation Learning (ROUSER) within the Information Bottleneck (IB) framework. This paper evaluates ROUSER on VRL generalization benchmark and conducts various experiments to demonstrate the advances of ROUSER.", "questions": "1.The six tasks in the DeepMind Control (DMC) suite, including \"reacher_easy,\" are commonly used to assess the sample efficiency of VRL. However, the DMC-GB also includes tasks like \"walker_stand\" for evaluating generalization capabilities. To provide a more comprehensive evaluation, you should run additional experiments on \"walker_stand\" to better assess your method's generalization performance in alignment with this benchmark.\n\n2.While CURL and DrQv2 have been evaluated on generalization benchmarks, they were not originally designed to address the specific challenges of generalization in VRL. To provide a more relevant comparison, you should include results for your method alongside established methods like SECANT [1] and PIE-G [2], which were specifically developed to enhance generalization in VRL. \n\n[1] Linxi Fan, Guanzhi Wang, De-An Huang, Zhiding Yu, Li Fei-Fei, Yuke Zhu, and Animashree Anand kumar. Secant: Self-expert cloning for zero-shot generalization of visual policies. In International Conference on Machine Learning, 2021.\n\n[2] Zhecheng Yuan, Zhengrong Xue, Bo Yuan, Xueqian Wang, Yi Wu, Yang Gao, and Huazhe Xu. Pre trained image encoder for generalizable visual reinforcement learning. In First Workshop on Pre training: Perspectives, Pitfalls, and Paths Forward at ICML 2022, 2022.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1.The method part of this paper has sufficient theoretical analysis and proof.\n\n2.The method is evaluated on various experiments.", "weaknesses": "See questions.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "6KdlKu0xVl", "reviewer": "Reviewer_238T", "summary": "This paper proposes robust action-value representation learning (ROUSER). ROUSER aims to acquire robust representation by encouraging the representation to capture Q-values while minimizing mutual information with state-action pairs to discard irrelevant features. ROUSER is applicable to different VRL algorithms and can improve performance in environments with visual distractions.", "questions": "NA", "limitations": "", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "- ROUSER aims to keep essential information by predicting the Q-value while neglecting irrelevant features by minimizing mutual information with state-action pairs.\n- ROUSER's objective is applicable to common VRL algorithms.", "weaknesses": "- The idea of incorporating Q-values to mitigate visual distractions is a commonly adopted design choice in VRL, e.g. [1,4]. I believe the effectiveness of ROUSER needs to be evaluated in a more standard setting with VRL generalization baselines.\n- To be specific:\n    - The [DMC-GB](https://github.com/nicklashansen/dmcontrol-generalization-benchmark) is a commonly adopted visual generalization benchmark with `color_hard` and `video_hard` setting [1,2]. I believe the evaluation results on DMC-GB would be more persuasive.\n    - For baselines, I believe comparing with generalization baselines is necessary. For example:\n        - TPC [3] proposes to regularize the representation to encode temporally predictable features only.\n        - [4] also proposes an information-bottleneck perspective to regularize representation.\n        - TIA [5] designs structured representation so that the task-relevant representation captures long-term rewards while the task-irrelevant one captures no reward information.\n        - PIEG [6] shows that representation from a pre-trained image encoder can boost the visual generalization significantly.\n\n\n[1] Look where you look! Saliency-guided Q-networks for generalization in visual Reinforcement Learning\n\n[2] Stabilizing Deep Q-Learning with ConvNets and Vision Transformers under Data Augmentation\n\n[3] Temporal Predictive Coding For Model-Based Planning In Latent Space\n\n[4] How to Learn Domain-Invariant Representations for Visual Reinforcement Learning - An Information-Theoretical Perspective\n\n[5] Learning Task Informed Abstractions\n\n[6] Pre-Trained Image Encoder for Generalizable Visual Reinforcement Learning", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "4sfbSzRni8", "reviewer": "Reviewer_uw2P", "summary": "The authors propose a novel temporal difference objective (ROUSE) for learning representations that capture long-term information are robust to distractions and perturbations for visual reinforcement learning. To make the proposed method robust, the authors utilize the information bottleneck principle, and make sure the learned representations of state action pairs captures the information about the reward, while minimizing the amount of information it captures about the state itself.\n\nROUSE is can be integrated with a variety of critic-based methods. The authors integrated ROUSE with DrQ, SRM and QRDQN, and test on DeepMind Control, MuJoCo, ProcGen, showing improved robustness compared to the baselines.", "questions": "Line 506: What information are you predicting in this frozen evaluation? Is it average of 300 rewards in the future? What exactly is the MSE between?\n\n\n##### Writing\n\n- 227: \"one of action-value representations\" : unclear wording, what does \"one of\" mean here?\n- 270: we support that a compressed reward representation... : \n- 299: model a reward model -> introduce a reward model\n- 364: supports -> allows\n- Sections 5.2 and 5.3 titles are not very informative. Applicability and effectiveness? what does that mean? Something more concrete would be better, like \"Extending ROUSER to tasks with discrete actions\" and \"Analyzing learned representations\"\n- Figure 4: please put the x and y axis labels in the figure itself, at least in small font, that would make it easier to read.-\n- 466: Reword to something like \"In this subsection, motivated by the result in theorem 4.4, we extend RoUSER to...\"\n- 502: all methods' prediction: what does that mean?\n- 504: \"use a 2 layer mlp with a learning rate of 1e4\". I understand what you're saying, but the sentence is awkward. I'd say \"we train a 2 layer mlp on top of frozen..., and use 1e-4 learning rate to train it.\"\n- 515: by searchings alpha... : reword this, this is not grammatically correct\n- 539: \"sheds light\" -- instead say \"we believe this method is general and applicable to other problems, such as self-driving\"", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "- ROUSE is general can be applied to any visual RL algorithm that learns a critic;\n- The proposed information-bottleneck objective is well-motivated mathematically. The authors provide detailed mathematical derivations of the objectives;\n- The method shows good results on the proposed benchmarks and improves robustness to color change and background video distractors;", "weaknesses": "- The paper is not clearly written, it took me some time to go through and understand it. I suggest improving the flow a bit. For example, something that could help is better structure. In the beginning of section 4, I would say something about the big picture of the paper, as \"We propose robust reward representations, then use td learning to learn robust value representations. Throughout, we use IB-based objective make representations robust.\" At the end of section 4, it would also help to list the objectives, or add the figure or pseudocode from the appendix. I understand you don't have much space, maybe decreasing the size of figures 1 and 2, or aggregating results over tasks could be a way to go here.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "uTIoEi0fy8", "reviewer": "Reviewer_6q1B", "summary": "This paper proposes a novel way for learning robust representations in image-based RL by learning through the proposed robust action-value representation using an information bottleneck framework. To address the challenge of unknown true action-values, this approach learns action-values indirectly through reward representations. The proposed method is validated in both continuous and discrete control domain. Experiments demonstrate improved performance over other image-based RL baselines on tasks with video background and object color changes.", "questions": "- In Section 5.3, the authors analyze the effectiveness of ROUSER in capturing long-term information with a fixed length of 300, It would be helpful to see how the effectiveness changes with varying lengths (e.g., 50, 100, 500).\n- The recent method MIIR [1] also considers learning robust representations using the information bottleneck framework and conducts experiments on the Distracting DMControl benchmark. Could the author discuss the differences between ROUSER and MIIR, such as their theoretical motivations and empirical performance comparisons?\n\n[1] Wang S, Wu Z, Wang J, Hu X, Lin Y, Lv K. How to Learn Domain-Invariant Representations for Visual Reinforcement Learning: An Information-Theoretical Perspective.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The paper is clearly written and easy to follow.\n- The motivation and explanation for learning robust action-value representations, and learning robust action-value representations via reward representations are presented clearly.\n- The experiment is conducted on both continuous and discrete control tasks.\n- Applying ROUSER on top of DrQv2 and SRM shows consistent improvement.", "weaknesses": "- Lack of training curves on Distracting DMControl benchmark for Figures 1 and 2: Including training curves would help readers understand the sample efficiency and stability of the proposed method.\n- In Figures 1 and 2, the baselines include both standard methods for image-based RL (e.g., DrQv2, CURL, TACO) and methods designed to be robust against distractors (e.g., SRM, CRESP, RAP). It would be better to reorganize the results, such as grouping methods by type or using different visual indicators for a clearer comparison. \n\nMinor: \n- Typo in line 430: Nore -> Note\n- In Figure 4, axis labels should be included directly on the plots rather than in the caption.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)