Self-supervised Color Generalization in Reinforcement Learning
Authors: Matthias Weissenbacher
Evangelos Routis
Yoshinobu Kawahara
Venue: Accepted by TMLR ()

Abstract: A challenge in reinforcement learning lies in effectively deploying trained policies to handle out-of-distribution data and environmental variations. Agents observing pixel-based image data are generally sensitive to background distractions and color changes. Commonly, color generalization is achieved through data augmentation. In contrast, we propose a color-invariant neural network layer that adopts distinct color symmetries in a self-supervised fashion. This allows for color sensitivity while achieving generalization. Our approach is based on dynamic-mode decomposition, which also accommodates spatial and temporal symmetries; we discuss the controlled breaking of the latter. We empirically evaluate our method in the Minigrid, Procgen, and DeepMind Control suites and find improved color sensitivity and generalisation.

Published in Transactions on Machine Learning Research (10/2024)
Self-supervised Color Generalization in Reinforcement Learn- ing
Matthias Weissenbacher
ubmatthias@gmail.co
Riken Center for Advonced Intelligence Project
Pyr-SAI Labs
Japan
Evangelos Routis
evongelos_rouiie routis Galumni brown. edu
Cousaly
London
United Kingdom
Yoshinobu Kawahara
Riken Center for Advanced Intelligence Project
Osaka University
Japan
Reviewed on OpenReview: https: openreview. net/ forum? id- 40mOPLRI88
Abstract
A challenge in reinforcement learning lies in effectively deploying trained policies to handle aut-of-distribution data and environmental variations. Agents observing pixel-based image data are generally sensitive to background distractions and color changes. Commonly, color generalization is achieved through data augmentation. In contrast, we propose a color- invariant neural network layer that adopts distinct color symmetries in a self-supervised fashion. This allows for color sensitivity while achieving generalization. Our approach is based on dynamic-mode dccomposition, which also accommodates spatial and temporal symmetries; discuss the controlled breaking of the latter. We empirically evaluate our method in the Minigrid, Procgen, and Deep.ind Control suites and find improved color sensitivity and generalisation.
1 Introduction
Reinforcement learning has seen tremendous success, but ensuring it works well on unfamiliar data still challenge. The original approach to generalisation in image-based RL, which used randomized image augmentations (Taskin et. al., 2020; Yarats et al., 2021b), was remarkably successful despite its brute- force nature. However, using image-based data augmentation can lead to overlregulaii to specific augmentations and thus may perform poorly when applied to certain environments. The more evolved methodology of selecting specific image augmentations in a self-supervised way leads to increased performance and alleviates some aspects of over-rogularisation (Raileanu et al., 2020; Hansen & Wang, 2021).
However, finding the balance between color sensitivity and color generalisation remains an open challenge, which is not addressed by the above-mentinn mentioned methods. Color sensitivity refers to the agent's ability to take action based on specific colors of objects in the environment, such as whether to stop or go at red or green traffic light, see Figure (1). In contrast, it is imperative that the agent's decision-making process remains invariant. to color changes in certain environmental factors, e.g., the color of houses, trees, or non-emergency vehicles in the context of autonomous vehicles.
An alternative data augmentation utilizing symmetrics in RL to improve generalization (Tang Ha, 2021; der Pol et al., 2020; Weissenbacher et al., 2022). Specifically, reinforcement learning may benefit. from

---

Published in Transactions on Machine Learning Research (10/2024)
both local and global symmetries, which preserve particular structure property within a neighborhood of point and throughout the entire space, respectively. In the above example, house or other vehicles are generally local, that is, they occupy a small area of the entire visual field of the agent. Enforcing local symmetrics through data augmentation becomes impractical as the sample size required grows exponentially with the number of image patches under consideration. This makes the approach highly inefficient and computationally expensive.
Our approach introduces the capability to develop self-supervised understanding of these local and global Layer (CiL) is based on the draaamienooddddoomsiitiioo
color symmetries. Our thcoretically incorporates spatial and temporal permutation invariance (PI). We do not make use of spatial PI in our model and break temporal symmetries in the case of frame-stacks as inputs, so that the time-wise order of frames is relevant.
which
In this work, we introduce a novel framework of self-supervised sclection of features. We apply the above- mentioned paradigm of symmetries to self-supervised selection of color features to replace image augmentating iin with random-conv and color-jitter filters. Consider the conventional training of an agent with randomly color-augmented images (see Figure (1). Depending the task, one could select subset of these images that are more suitable. Our method effectively performs this selection of 'useful' color augmentations in a self-supervised manner. CiL. admits different inductive biases categories (T)-(III) admitting an increased amount of trainable parameters corresponding to specific subsets of augmented images. The most suitable bias selected during training in a self-supervised fashion.
In section 3 we introduce both the architecture for the CiT. layer well as provide formal propositions which validate its symmetry properties. In section 4.1, we show that in the simple example of adding safe blue river (water) to the mini-grid LavaCrossing environment data augmentation random-cony and color-jitter is detrimental to learning, while CiL is able to navigate the selsivvensitvve environment. Moreover, we find improved color generalisation compared to the baseline. In section 4.2 we perform our main experiments on the Procgen environment. The code is made public at GitHub.
2 Background
2.1 Reinforcement Learning& & Symmetries
Reinforcement Learning. A Markov Decision Process (MDP) is mathematical framework for modeling decision-making problems in stochastic environments. MDPs are characterized by a tuple AS, P, R, 7), where S is a finite set of states, A is a finite set of actions, P is the transition probability function, R is the reward function, and y E 10 is the discount factor. In RL, one aims to learn optimal decision- making policies in MDPs. A policy, denoted as क S + A, is a mapping from states to actions. The optimal policy n" maximizes the expected cumulative discounted reward, given by the value function (s)(s) EE 7 7'R(s,, | so = s,a, ~ #{ [sr]], where the expectation is taken over the sequence of states and actions encountered by following the policy x. The optimal policy x" is the one that satisfies 1" (s) 2 (s) for all E S and any other policy π.
Symmetry groups. Symmetry transformations are special cases of mathematical groups, i.e. sets that contain the identity element and are closed under multiplication and inverses. In this work we encounter the group of orthogonal matrices, in particular the three-dimensional rotation group SO3, i.e. the special orthogonal group in dimension 3. Any element of the latter can be written in terms of three independent parameters i.e. the rotation angles along the x,y and z-axis, respectively, called the Euler angles.
Invariance. Invariance is a foundational concept in understanding how functions respond to symmetries of their inputs. Before defining these concepts, we introduce some notation. Let f be function that maps elements from space S to space S'. Let. G be a symmelry group acting on S and let us denote by an individual transformation in G. We shall denote the action of the transformation on an element s of by g 8. A function f is invariant with respect to a set of transformations (symmetry group G) the application of any transformation from this set its input does not change the function's 's output. Mathenatically, this is expressed as f(g 8) f(s) for every transformation g E G.
2

---

Published in Transactions on Machine Learning Research (10/2024)
Training
Unseen
Generalisation Sensitivity
Ganenilissiion
(Cole-Serrettttty
A
A
Augmentations
OURS: Self supervised Symmetries (Unseen)
Geseralataion
Symmetry Mendilcations
Cokoon Bersiliity
Observation
Figure 1: Top: Overview of the balance of color generalisation color sensitivity required for RT tasks, in the example of the Jumper environment of the Procgen suite. Bottom: Conventional random color-augmentation i.e. color-jitter (left.) compared OUR self-supervised learning of color-symmetries (right), in the example of Category II color symmetries - horizontal axis- per each feature of the model vertical axis- The vertical axis shows different features- wise choices of Category II symmetries. The horizontal axis shows identical images under the particular Category II symmelry. In other words, CiL produces the same feature-wise output if the input image varies according horizontal axis.
2.2 Dynamic Mode Representation
The dynamni-nood--eeommeitton (DMD) algorithm fits a linear operator IC that advances the state of a system, 8 IR" forward in time (Schmid, 2010; Tu et al., 2013; Kutz et al., 2016; Rowley et al., 2009; Williams et al., 2015). Thus, the linear dynamical system given by
St,-1 Kst, for == N
The operator K. is an approximation of the Koopman operator restricted to observables given by direct measurements of the state s. In practice the operator K. is computed from a collection of snapshot pairs of the system {St.. St, h= In principle, for DMD the times need not be sequential or evenly spaced, however this is the case for most RT. settings. These snapshots are arranged into two data matrices, S and as
S= ( 81g .... Str) and S' = (8t2) Sts ..., 8tx18))
thus equation 1 may be rewritten as S ~ K.S. The computation of operator K. from given transition dynamics data can be reformulated as an optimization problem
K arg min S K Klp St
K
where | |F is the Frobenius norm and i denotes the pscudo-inverse. The pscudo-inverse may be computed by applying singular value decomposition (SVD). In the case of image data with color features we stack them in the time-direction as
S (St, St1 St. Sex Sty: Stx)
S' (8tg, Stgs Stg3 ... St 8t 8tg+1)
(4)
3

---

Published in Transactions on Machine Learning Research (10/2024)
Category
Fermeed Colp Symeetty por Twees
Sceeilo
Catagory
Speeile Colee Symmaryy Feadie
Figure 2: Three categories of the feature-wise color symmetries of CiL in the example of the Jumper enviromment of the Procgen Suite (Cobbe et al., 2020).
A time sequence of data is very commonly modified by invertible linear transformation acting on each state as well as an invertible linear transformation acting on the time-series as
S
S 5C
(5)
Vi = 1.....N and where R Raxn and 3 We restrict ourselves to orthogonal transformations i.e. R- R" and --1=""
2.2.1 Leading Rank Representation
Since the matrix operator IC typically admits n² elements, for nigh-dimensional data, it is intractable to represent this operator. Instead, one may compute its leading spectral decomposition. In particular, for time series pixel data, the sequence length N is much shorter (in the context of RL) than the number of pixels in the 2D image, thus the matrices S and S' have far more rows than columns. It follows that the size of the set of eigenvectors of K corresponding to its nonzero eigenvalues is at most N times the number of colors (input channels of cach image). In practice, the effective rank r of the data matrices S and S' and hence the operator K. generically is even lower.
The algorithm of the leading order rank approximation (rDMD) can be outlined as follows. Instead of computing K in equation 3, we may project the latter onto its first r singular vectors. With U,, E,, V, the rank-r restricted singular value decomposition S one can approximate the pseudo-inverse as S x U, E, V Then the operator Ic, can be defined
rDMD: K- = U, S' V
The leading spectral decomposition of K. may be approximated from the spectral decomposition of the much smaller Xr, as IC, W = WA The diagonal matrix A contains the DMD eigenvalues, which correspond to eigenvalues of the digh-dimensional matrix K.
3 An Adaptive Color Invariant Layer
In this section, we delve into the architecture of our color-invariant layer (CiL), offering a detailed technical overview. Subscquently, in section 3.1, we provide its thcorctical underpinning. Central to CiL the strategy of partitioning the image into patches, drawing parallels with vision transformers (Dosovitskiy et al., 2020). For each of these patches, we then compute the Singular Value Decomposition (SVD), in some cases, the mean over adjacent patches, as depicted in Figure (3). This approach is twofold: firstly, it minimizes the computational demands of SVD; secondly, it paves the way for achieving local color invariance.
Conceptually, we divide color symmetries, both local and global, into three distinct mathematical categories (II) Very general invariance (I); (II) the model can self-supervised way learn a specific color spectrum
The eigenvectors of K,, are not the same as the eigenvectors K- there is precise relation between the two (Brunton, 2019)
4

---

Published in Transactions on Machine Learning Research (10/2024)
per each 3-chamnel for which it is invariant, see Figure (1) and Figure (2); for type (III) the model in self-supervised fashion learns a pair of color-transformation which it understands to be the same, ie is invariant under- per each 3-channel -. Category (1) is comparable to random-cony and color-jitter filters in that it not color-sensitive. Categorics (II)-(III) increasingly trade off gencralisation to gain more color sensitivily.
For simplicity, let us discuss the single frame case first i.e. 1 for which the core functionality of the CiL layer is given by
CiL S
S W V
where V is computed by the SVD of S. and W is 3x3 matrix of trainable weights. 3 The choice of the latter results the following categories:
Category I
W the identity matrix without any trainable parameters, equation 7 is invariant.
under general orthogonal transformations.
For We E SO3, i.e. in the general rotations group equation 7 is invariant under a continuous one- one-parameter family of color transformations. For a given element in SO3 corresponds to a rotation in threo-dimensional color space and the layer is symmetric w.r.t. any other oration along the same axis. The specific axis chosen by setting the three Euler angles i.e. independent. parameters in W.
Category II:
Category : For W a symmetric (real) 3 x matrix, equation is invariant under a set of (or 8) specific color transformations which are chosen by setting the parameters in W.
For visualisation Figure 2. For multiple frames Lee N >>> need to take into account the time-series nature according to equation 4.
CiL-Stack S + Syy Trr'T" Wi ((S))
(8)
where for notional simplicity we split S' color and stack dimension, and we have used Einstein summation
convention with i,j, ,, and V(S) is computed via the SVD. We introduce trainable weights T which are unconstrained an break the time-reversal symmetry of the expression equation 8; T, ,,T', I...., N.
3.1 Theoretical Aspects
Proposition (CiL Symmetry Invariance). Our approach equation invorriant under:
right matriz multiplication by orthogonal L equation W 1 i.e. equal to the identity matria Category T-
by a one- ine-porameter family of (special) orthogonal Ls equation 5 if W E SO(3). The specific depend on the details of weights in W, see the appendis section (B.1) and (C.2.1) for exact relations Category 11
by 3 (6 2) symmetric L1, L2, La equalion 5 isf W is symmelric. The specific C depend on the delails of weights in W, see the proof section (B.1) and (C.2.1 for exact relations Category III
Proposition 2 (rDMD Symmetry Invariance). The rDMD approach equation 6 is invariant under left and right matriz multiplication by orthogonal C, equation 5, respectively.
2e refer to 3-charanel the input color-channcls which doplicated in lager. As for color-channcls V and mostly identical them synonymously from here 4The expression in the parenthesis in proposition (1) refer the case when normalise by determinant. In the default use V/det(V).
5

---

Published in Transactions on Machine Learning Research (10/2024)
Localty Culor Symmetry inverient under Category E or respecttr
Recerstivd iuge von Putdees
Weighta of Category & Br
Duplicate Patch & Vor.a
SVD
8VD
BVD
IVV
SVD
INPUT
For ench of the IP patches sttechh the neighbouring patches within a cortain window sidc
moan cria
moan
Divided into patchen
porches
parcches
Paira of Neighbouthood and Center Putch
Figure Architecture diagram of CiL for N no frame stacking.
Controlling spatial and temporal symmetries of rDMD. Among other symmetries, equation 3 and equation 6 admit permutation invariance and general rotation of pixels as long the operation acts each time step in the frame stack of input images alike, see Figure (4) (c) and (d). It is evident that those invariances admit no practical use case. It is this desirable to remove those symmetries while keeping the useful ones e.g. Figure (4) (b). In (Weissenbacher et al., 2024) recently a graph symmetric approach to symmetry breaking in pixel data was discussed. We define
Kbreak = KU = U,
(9)
where G is the graph matrix discussed in (Weissenbacher et al., 2024) and is such that [9, R = G R
Oie. they commute if R is a simple left right rotation or flip of the patch. In particular, for general permutation of pixels in the patch they do not commute. The matrices g RRaxn, E NRN and W E R contain trainable weights. While we do not directly use equation involving graph matrix in this work our formalism can be extended to include the latter.
Proposition 3 (rDMD- Spatial Symmetry Breaking). Eq. (9) is invariant under left and right matrix multiplication by orthogonal L, R equation 5, which commate with G. Be G E R² x R is given by the symmetric graph-matrir, i.e. entries have shared weights if the distance between patches is the same. Then the symmetries of rotations, flips are preserved while the symmetry of undesirable general pizel permutations and orthogonal transformations is broken.
Proposition 4 (CiL-Stack Imposing Time-order). By adding the weights T E RNNN version of CiL equation one imposes a time-ordering of the fromes, Le. one generally breaks the orthogonal symmetries applied to the dack-dimensson.
of the frame slack
On the importance reducing an abundance of symmetry In the previous section we have laid the theoretical basis of how specific symmetries of CiL may be broken. In this concluding section we informally highlight some of the findings of section 3.1. In Figure (4) we illustrate patch-wise and global transformation of spatial and color. We proof in the appendix C.2.4 that CiL is not invraint under regional color changes Figure (4) (a). This property is crucial to be able identify color sensitive feature e.g. teh red green traffic light.
3.2

---

Published in Transactions on Machine Learning Research (10/2024)
f
A
h
(a) Regional color changes. Broken by CiL, appendix (C.2.4).
(b) Orthogonal color changes i.e. application of L.
(c) Patchwise pixel per- mutations and blurrs (spa- tial transformation).
(d) Patchwise orthogonal
global
matrix multiplication by
R.
Comparison of symmetries which aid learning (a) and (b) contrasted to those impeding the ability to learn crucial features (c) and (d).
Figure 4:
Seen Training
Water
0.8
CNN
0.7
catagory
Sale Puth
Farst Pam
0.6
category
color-jitter
random-come
0.4
0.2
Agent
Dod
0.0
1e7
(a) Modified Lavacrossing, River (blue) are safe passage through Lava stream (orange).
(b) Training rewards, averaged over 4 random seeds.
Figure Modified lava-crossing environment overview and empirical evaluation of CiL. Comparison of CiL Color-jitter and Random-conv data augmentation, the latter are detrimental for learning.
In Figure (4) (b) global taloforransooom by distinct symmetric matrices L are shown. Global refer to the choice of the local neighborhood of SVD in Figure (3) to be the entire image, i.e. the same V is the shared across all patches.
In Figure (4) (c) and (d) we highlight patch-wise spatial transformation e.g. pixel-permutations and general orthogonal transformation. Those result. apparent non-useful symmetry properties. A model with those too abundant symmetrics fails learn relevant fcatures necessary for the agent to navigate the environment.
In Figure 3, the different patch colors for illustration only. The patches are not independent; the same color invariance is learned for each feature across all patches, CiL's weights are shared among them. One might wonder if patch-wise color augmentation could yicld similar results to CiL. However, empirical tests on Minigrid show that patch-wise augmentations, like global color augmentations, also fail to learn elfectively. Additionally, to achieve the same invariance properties as CiL, the model would require an impractical number of patch-wise augmentations, namely: (#local color augmentations) patches
4 Empirical evaluation
4.1 Modified Minigird Lavacrossing Test sensitvvns & Generalisation
The LavaCrossing environment, standard in the MiniGrid toolkit (Chevalier-Biiseer et al., 2019), requires agents navigate a goal (green square) without falling into lava (orange squares). We modify the environment by adding:

---

Published in Transactions on Machine Learning Research (10/2024)
Seen / Training with Berries & Crops
Agent gains reward by picking up berries or crops
Berries (reward 0.4)
Crops (roward 0.15)
Figure Modified Lava-crossing with additional option for the agent to collect rewards by harvesting berries (purple) and crops (green).
1. safe river (water) to the environment which constitutes the only safe pathway through the lava stream, see Figure (5a). Thus the agent must learn to take "swim" in the river which is a color sensitive choice as the lava, needs to be avoided.
optionally to point (1) we introduce berry-fields and crop-fields see Figure (6) upon collecting them the agent receives additional rewards.
Thus depending on the colors blue and orange of river vs. lava (purple and green of berries and crops) see Figure (6) opposing actions need to be taken by the agent to successfully reach the green goal. In other words, this modified environment an ideal test-ground for an agent's ability of developing color-sensitivity, see Figure (10).
Although MiniGrid setups are usually partially observable, we adjust ours for full observability and increase the default observation size from 9 to 14 x 14 pixels. In particular, render the environment and subsequently down-scale it. The experiments with deep Q-learning (IMPALA (Espeholt et al., 2018)) focus on difficulty level of our modified Lava-crossing environment.
We evaluated the CNN baseline against configurations using color-jitter or random-cony data augmentation. These methods were unsuccessful in learning the desired behavior as they hindered the agent's capacity for color-sensitive makiii..-making categories T, II, and III, learning proceeded without obstruction. Refer to Figure (5b) for the training rewards. Our findings suggest that not only does CiL achieve inherent generalization, but it also equips the model with color-sensitive decision-making skills.
contrast, when we incorporated CiL with the CNN for color-symmetry
Furthermore, Figure (5b) and (7a) illustrates the enhanced performance and sample efficiency of CiT. + CNN, for both environment versions, respectively. Morcover, Figure (7b) shows improved color gencralisation of CiL. Notably, this improvement observed without any adjustments to the hyper-pann paremeersseeerss of the RL algorithm.
Next, extend the empirical results our custom LavaCrossing environment and present evaluations of point (2) below. We modily the e by adding we introduce berry-fields as well crop-fields upon collecting them the agent receives additional rewards.
In evaluating the performance CiL on the environment option (2), which inchudes both berries and crops, we observe results that are consistent with those obtained in environment option (1). This can be seen Figure (7a), which presents a visual representation of the findings.
The LavaCrossing environment, standard in the MiniGrid toolkit "Chevaller-Bolvver et al., 2019), requires agents navigate to goal (green square) without falling into lava (orange squares).

---

Published in Transactions on Machine Learning Research (10/2024)
CNV
cenepety 1x11
caleggty
Letoggiyy
Lategarr
12
catngTte
categrry
1.0
32-
0.4
(b) Color Generalisation: rewards, averaged random seeds and evaluation episodes; plot the standard bars, respectively. Augmentation parameter are from right as Figure
00 3u7
(a) Training rewards, averaged over random seeds.
Figure 7: Empirical evaluation of CiL modified lava-crossing environment (2) with berries and crops. Comparison of CiT, to color-jitter and random-convolution data-augmentation, the latter are detrimental learning- Comparison of Categories 11, with simultaneously I -1 Both Figure 7a and 7b show mean and standard error.
as the input of the CNN I&/1 and batch-coneattnntted i.. used
Figure 8: Modified Lava-crossing color variations for Figure visualizing augmentation parameter from left right [0, L88]
When lata-auumentttion is applied, we notice a detrimental effect on the training process. This in line with our observations from environment option (1), suggesting that dala-augmentation might not be beneficial for this specific task.
Furthermore, we observe that the CiL category II demonstrates superior performance when compared to the CNN baseline when no data augmentation is applied. This performance advantage more pronounced in environment option (2) than what was observed environment option (1), as depicted in Figure (5b). This indicates that the benefits of using CiL category become more evident when the complexity of the environment increases, showcasing its robustness and effectiveness in more challenging scenarios.
We evaluate our trained models color varied version of the environment, Figure (8); comparing CNN to CNN + CiL for categories I,I1, and II, respectively, We find improved generalisation of CiL, see Figure (7b).
4.2 Procgen Benchmark & Distracted Deepmind Control suite
We evaluate our model on both the Procgen gencralization benchmark (Cobbe et al., 2020) which consists of 16 procedurally generated environments with visual observations; well as the Deepmind Control suite (DMControl) (Tassa et al., 2018) with additional visual background video distraction (Hansen & Wang, 2021).
The Procgen benchmark consists of sixteen procedurally generated games. Each game corresponds to distribution of partially observable Markov decision processes (POMDPs) q(m), and each level of a game corresponds to a POMDP sampled from that game's distribution m ~ q. The POMDP m determined by the seed (i.c. integer) used to generate the corresponding level. Following the setup from (Cobbe al., 2020), agents are trained on fixed set of n 200 levels (generated using seeds from 1 200) and tested the full distribution of levels (generated by sampling seeds uniformly random from all computer

---

Published in Transactions on Machine Learning Research (10/2024)
Comparison between CE, and Rand-FM
Algorthmm
80
DVAC cop (CMN+CL) Rand-FN (CWN)
Percettage
StarPAct
FrutBot
BigFsh Eivironment
Figure 9: Comparison of CiL random-conv augmentation top of PPO backbone.
integers). We evaluate method with IPO/DrAC (Schulman et al., 2017; Raileanu et al., 2020) with added crop augmentation.
Moreover, we evaluate our method with SAC (Haarnoja et al., 2018) with standard hyper-parameter settings, and simply add CiL as a preliminary layer. To evaluate generalization of our method and baselines, we test methods under challenging distribution shifts from by adding background videos (Hansen & Wang, 2021). We provide a proof-of-concept for CiL's ability generalize to background videos in DMControl. This successful test is significant us, SVD/DMD can be highly sensitive to dynamic changes introduced by background videos. The DMControl test demonstrates CiL's effective generalization to new data distributions relevant for real-world tasks, rather than aiming to outperform statc-of-the-art (SOTA) methods.
Based on the Procgen and DM-Control (DMC) results in Tables (1) and (3), and the Procgen ablations in Table (2), we can draw three main conclusions
First, CiL can scale to Procgen and DMC tasks.
Second, CiL may be simply added as a preliminary layer to existing architectures; across variety of domains and tasks and does not require any seper-parameted tuning to achieve comparable or improved generalisation performance. The performance gains are significant for environments where color-scnsitivity is required c.g. StarPilot and BigFish, in particular compared rand-conv and color-jilter data augmentations.
Third, the ablation study of CiL on Procgen finds that its performance is consistent across wide range settings.
In Figure (9), we juxtapose CiL with random-conv augmentations. Specifically for the Starpilot and Bigfish environments, the use of random-conv (Rand-FM) hinders learning. In contrast, CiL offers advantages over both PPO and DrAC+ Crop, especially in these environments. This aligns with main results Table 1 where random-cony (as well as color-jitter) are added to DrAC. The training curves show comparable sample-effcieney of CiL, see Figure (11).
The primary conclusion the paper is that mixture different categories yields the best performance. The general principle is that by incorporating all categorics-cact with distinct inductive biases- the network is better equipped lo select the most appropriate features.
4.3 Limitations
CiL is CPU extensive by employing SVD as existing GPU implementation of SVD barely bring speed-ups, and rather lead to bottle-necks. Repeating input data (refer to Figure (3)) assign different symmetry weights across channels can be demanding terms of memory and computation. Particularly, using SVD locally in CiL. increases runtime by factor of 3-5 compared to the CNN baseline. This work serves
10

---

Published in Transactions on Machine Learning Research (10/2024)
Table 1: We experiment on the Procgen environments. compare to the PPO baseline and Rand-FM (Lee al., 2020) with ResNet architecture. We report min-max normalised scores averaged over the testing steps 23M-25M, and seeds, respectively. We use DrAC + crop augmentation for our model which we use preliminary CiL and the same subsequent ResNet. The CiL used local neighbourhood of 5 patches both vertically and horizontally and categories I, II and III. Both random convolution and color-jittes augmentations are added (for Conv&JJtter with equal likelihood). The standard error for CNN+CiL, well baseline is < 0.2% for each game individually.
Procgen
Algorithm
Model
StarPilot
FruitBot
BigFish
Coinrun
Jumper
Average
PPO
CNN
35.6%
85.9%
7.8%
65.4%
53.3%
49.6%
DrAC (Crop)
CNN
50.0%
86.5%
21.6 %
64.5%
52.4%
55.0%
DrAC
CNN
47.3%
84.8%
14.2%
64.9%
56.5%
53.5%
DrAC (Crop+ ColorJitter)
CNN
52.2%
85.9%
13.1%
64.9%
53.8%
54.1%
DrAC (Crop + Conv&JJiterr
CNN
41.8%
84.0%
9.1%
65.0%
56.2%
51.9%
DrAC (Crop)
CNN+CIL
55.4%
86.9%
26.6%
64.11%
53.4%
57.3%
Rand-FM
CNN
10.1%
79.7%
-1.1%
71.2%
47.8%
41.5%
Table 2: Procgen ablation study comparing categories I, TI, and TIT as well as different local neighborhood sizes. The superscript refers the weight categories used in CiT. The subscript 1; 3, 5, refer to the size of the local neighborhood used to compute the SVD; where denotes the global CiL, i.e. the mean is computed over patches of the entire image.
Procgen Task
CiL!
CiL!&
CiL!
CiL
CiL CiLg-1"!
StarPilot
53.0%
44.8%
47.8%
51.2%
50.7%
55.4%
FruitBot
86.6%
86.3%
88.0%
86.7%
87.6%
86.9%
BigFiish
21.6%
14.0%
21.3%
11.7%
23.7%
26.6%
Coinrun 63.5%
60.8%
53.7%
64.2%
64.9%
64.4%
Jumper
53.0%
46.6%
48.3%
54.7%
54.4%
53.4%
Average
35.5%
50.5%
51.8%
53.7%
56.3%
57.3%
Table 3: DM-control with background video distraction as well as color changes. We evaluate the trained models after 100k steps over episodes, respectively, and report the mean and standard deviation over seeds.
DMControl
Domain-Task
Distraction
CNN
CiLÇ&l kll
video-rasy
493.1
125.0
497.9
118.8
Walker-Walk
video-hard
137.5 + 72.6
112.6 +: 48.1
video-casy
167.3 1 99.9
169.2 1: 58.5
Cheetah-Ron
video-hard
41.1 = 26.1
61.8 + 44.6
proof of principle, highlighting a new rescarch direction with scope for future technical enhancements. We effectively reduced overhead by incorporating the SVD of the input frame-stack into the replay-bulfer, demonstrated in our DMC experiments.
5 Related Work
Symmetry is prevalent implicit approach in deep learning for designing neural networks with established equivariances and invariances. The literature symmetries in Vision Transformers (ViTs) (Fuchs et al., 2020; Romero & Cordonnier, 2021) is relatively limited compared to CNNs (Zhang & Sejnowski, 1988; LeCun et al., 1989; Zhang, 1990), recurrent neural nerssrks (Rumelhart et al., 1986; Hochreiter & Schmidhuber, 1997), graph neural networks (Maron et al., 2019; Satorras et al., 2021), and capsule networks (Sabour et al., 2017). Permutation invariance ViTs and attention mechanisms has been examined in (Loe et al., 2019), demonstrating improved out-of-dee of-diitribtttion generalization in RL from pixel data (Tang & Ha, 2021). CiL may be added top of ViTs well as CNNs.
11

---

Published in Transactions on Machine Learning Research (10/2024)
Conventionally sample efficiency is enhanced by data augmentation (Krizhevsky al., 2012). Simple image augmentations, such as random crop (Laskin et al., 2020) or shift (Yarats et al., 2021a; 2022), can improve RL generalisation performance; in particular when combined with contrastive learning (Agarwal et al., 2021). CiL is complementary to data-angmentation.
3D rotations, described by the 50(3) group, have found their place in various areas of vision-based AT tasks, however applied on the spatial and not color-features. From 3D medical imaging (Vieweg al., 2015) over 3D object recognition (Qi et al., 2017; Coors et al., 2018) transferring neural networks from simulations to real-world robotics (Tobin et al., 2017). Furthermore, deep learning on Lie Groups, encompassing SO(3), has evolved with architectures like the Spherical CNNs (Cohen et al., 2018), inherently managing rotations.
Symmelry-based representation learning (Balaraman & Andrew, 2004; van der Pol et al., 2020; Higgins et al., 2018) refers to the study of symmetries of the environment manifested in the latent representation and was extended to environmental interactions in (Casellos-Dupré et al., 2019). In (Weissenbacher al., 2022), symmetries of the dynamics are inferred in a self-supervised manner. More recent works equivariant. methods in RL include (van der Pol & Welling, 2019; Wang & Walters, 2022). These approaches mostly complimentary to CiL.
6 Conclusions
In the context of reinforcement learning, introduced CiL, a neural network layer capable of self-supervised adaptation to various color symmetries. Empirical evaluations show that CiL outperforms both random-conv and color-jitter data augmentations in environments where color sensitivity is pivotal. Furthermore, CiL maintains performance in other environments without necessitating hyperparameter tuning for the associated RL algorithms.
Impact Statement
This paper presents work whose goal is to advance the field of Machine Learning. There are many potential societal consoquences of our work. In particular, while improved color-sensitivity while maintaining color- generalisation is desirable for many lasks there are risks, such as an agent learning to make decisions based on an individual's skin color. However, we conclude that such a behavior would likely arise from biases in the training dataset or environment, rather than from our method itself.
Acknowledgments
Matthias Weissenbacher would like express his gratitude to Rishabh Agarwal (Deepmind) for initial collabo- ration and helpful comments on the draft. And, we would like thank Y. Nishimura for technical support. This work supported by JSPS KAKENHI Grant Number JP22H00516 and JP22H05106 JST CREST Grant Number JPMJCR1913.
References
Rishabh Agarwal, Marlos Machado, Pablo Samuel Castro, and Marc G Bellemare. Contrastive behavioral similarity embeddings for generalization in reinforcement learning. arXiv preprint 1rXiv:2101.05265, 2021.
Ravindran Balaraman and G. Barto Andrew. Approximate homomorphisms: framework for non-exact minimization markov decision processes. In International Conference on Knowledge Based Computer Systems, 2004.
Steven L. Brunton. Notes on koopman operator theory. 2019. URL 229348032.
Hugo Caselles- Dupré, Michael Garcia Ortiz, and David Filliat. Symmotry-based disentangled representation learning requires interaction with environments. In H. Wallach, H. Larochelle, A. Beygelzimer, F. d' d'Alch
12

---

Published in Transactions on Machine Learning Research (10/2024)
Buc, E. Fox, and R. Garnett (eds.), Advances in Neural Information Processing Systems, volume 32. Curran Associates, Inc., 2019.
Maxime Chevalier-Boisee Damitry Bahdanau, Salem Lahlou, Lucas Willems, Chitwan Saharia, Thien Huu Nguyen, and Rémi Lebret. Babyai: First steps towards grounded language learning with human in the loop. International Conference Learning Representations, 2019.
Karl Cobbe, Chris Hesse, Jacob Hilton, and John Schulman. Leveraging procedural generation to benchmark reinforcement learning. In Hal Daumé III and Aarti Singh (eds.), Proceedings of the 37th International Conference on Machine Learning, volume 119 of Proceedings of Machine Learning Research, pp. 2048 2056. PMLR, 13 18 Jul 2020.
Taco S. Cohen, Mario Geiger, Jonas Köhler, and Max Welling, Spherical cnns. In International Conference on Learning Represenlations (ICLR), 2018.
Benjamin Coors, Martin Roser, and Andreas Geiger. Spherenet: Learning spherical reprosentations for detection and classification in omnidirectional images. In Proceedings of the European Conference Compuler Vision (ECCV), 2018.
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszkoreit, and Neil Houlsby. An image is worth 16x16 words: Transformers for image recognition at scale, 2020.
Lasse Espeholt, Hubert Soyer, Remi Munos, Karen Simonyan, Volodymyr Mnih, Tom Ward, Yotam Doron, Vlad Firoiu, Tim Harley, lain Dumning, et al. Impala: Scalable distributed deep-rl with importance weighted actor-learner architectures. arXie preprinl orXiv:1802. 1502:0000..161., 2018.
Fabian Fuchs, Daniel Worrall, Volker Fischer, and Max Welling. transstrrassoommes:: 3d roto-translation equivariant attention networks. In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin (eds.), Advances in Newral Information Processing Systems, volume 33, pp. 1970-1960 1981. Cur- ran Associates, Inc., 2020. URL 11221177eeba27911112222255553333 Paper.. pdf.
hegpes.nuurppp.. cc/paper filles/paper//22//EEL
Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with a stochastic actor. In Jennifer Dy and Andreas Krause (eds.), Proceedings of the 35th International Conference on Machine Learning, volume 80 of Proceedings of Machine Learning Research, pp. 1861 1870. PMLR, 10 Jul 2018
Nicklas Hansen and Xiaolong Wang. Generalization in reinforcement learning by soft data augmentation. In International Conference on Robotics and Automation, 2021.
Irina Higgins, David Amos, David Pfau, Sebastien Racaniere, Loic Matthey, Danilo Rezende, and Alexander Lerchner. Towards a definition of disentangled representations, 2018.
Sepp Hochreiter and Jürgen Schmidhuber. Long short-term memory. Neural computation, 9(8):1735 1780, 1997.
Zhengyao Jiang, Pasquale Minervin, Minqi Jiang, and Tim Rocktäschel. Grid-to-graph: Flexible spatial relational inductive biases for reinforcement learning. In AAMAS 2021, 2021.
Alex Krizhevsky, Iya Sutskever, and Geoffrey Hinton. Imagenet classification with deep convolutional neural networks. In Advances in Neural Information Processing Systems, pp. 1097-1105, 2012.
Nathan Kutz, Xing Fu, and Steven Brunton. Dynamic mode dccomposition for big data: A reduced-order model with interpolative decomposition. SIAM Journal on Applied Dynamical Systems, 15(2):1327-1344 2016.
13

---

Published in Transactions on Machine Learning Research (10/2024)
Misha Laskin, Kimin Lee, Adam Stooke, Lerrel Pinto, Pieter Abbeel, and Aravind Srinivas. Rein- forcement. learning with augmented data. In H. Tarochelle, M. Ranzato, R. Hadsell, M.F. Bal- can, and H. Lin (eds.), Advances in Neural Information Processing Systems, volume 33, pp. 19884 19895. Curran Associates, Inc., 2020. URL https://proceedingg newrine httppp/p//ppppe/////11111 hww.eaaaba411111aaee22aaaaaaaaa Pappee Paf
Yann LeCun, Bernhard E Boser, John S Denker, Donnie Henderson, Richard E Howard, Wayne Hubbard, and Lawrence D Jackel. Backpropagation applied to handwritten zip code recognition. Newral computation, 1(4):541-551, 1989.
Juho Lee, Yoonho Lee, Jungtack Kim, Adam Kosiorck, Scungjin Choi, and Yee Whye Tch. Set transformer: A framework for attention-basee permutation-invaaaan neural networks. In Proceedings of the 36th International Conference Machine Learning. PMTR, 2019. URL https://proceedling. mlr. mlr. wtw.ggesss..rrr.ss/////. v97 v97/1ee19d.ho htmll.
Kimin Lee, Kibok Lee, Jinwoo Shin, and Honglak Loc. Network randomization: simple technique for generalization in deep reinforcement learning. In International Conference on Learning Representations, 2020. URL https://openrer// htw.pppawweeeeeeetetooorrrea.eae .et.://orrur.....CCg...B..
Haggai Maron, Heli Ben-Hamu, Nadav Shamir, and Yaron Lipman. Invariant and equivariant graph networks. In International Conference on Learning Representations, 2019.
Charles R. Qi et al. Pointnet: Deep learning on point sets for 3d classification and segmentation, In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition (CVPR), 2017.
Roberta Railcanu, Max Goldstein, Denis Yarats, Ilya Kostrikov, and Rob Fergus. Automatic data augmentation for generalization in deep reinforcement learning. arXiv preprint urXin:2006.12860 2020.
David W. Romero and Jean-Baptiste Cordonnier. Group equivariant stand-alone self-attention for vision. International Conference Learning Represendations, 2021. URL http:://pppneereeeeeeee/ttrtoroe Jn06GM.
Clarence W Rowley, Igor Mezić, Shervin Bagheri, Philipp Schlatter, and Dan S Henningson. Spectral analysis nonlincar flows. Journal of fluid mechanics, 641:115-127, 2009.
David E Rumelhart, Geoffrey E Hinton, and Ronald Williams. Learning representations by back-propagating errors. Nature, 623(6088))5533 (038)))333-3366666 1986.
Sara Sabour, Nicholas Frosst, and Geoffrey Hinton. Dynamic routing between capsules. In Advances in Neural Information Processing Systems, pp. 1856-38666 2017.
Victor Garcia Satorras, Emiel Hoogeboom, and Max Welling. E(uivequuvariant graph neural networks. In International Conference on Learning Representations, 2021.
Peter Schmid. Dynamic mode decomposition of numerical and experimental data. Journal of fluid mechanics, 656:5 28, 2010.
John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms, 2017.
G. Strang. Introduction to Linear Algebra. Wellesley-Cambriige Press, Wellesley, MA, edition, 2003.
Yujin Tang and David Ha. The sensory neuron transformer: Permutationniivaaraae neural networks for reinforcement learning. In A. Beygelzimer, Y. Dauphin, P. Tiang, and J. Wortman Vaughan (eds.), Advances in Neural Information Processing Systems, 2021. URL https:// http:r/wwwee.eee//orou////dddeee wet//forumm.devtI--A-me
Yuval Tassa, Yotam Doron, Alistair Muldal, Tom Erez, Yazhe Li, Diego de Las Casas, David Budden, Abbas Abdolmaleki, Josh Merel, Andrew Lefrancq. Timothy Lillicrap, and Martin A. Riedmiller. Deepmind control suite. CoRR, abs/1801.000 00690, 2018. URL http://arxivv www.aiiiv./......e....eeee wttp://o////..oooeeeeeeeoo
14

---

Published in Transactions on Machine Learning Research (10/2024)
Josh Tobin et al. Domain randomization for transferring deep neural networks from simulation to the real world. In Proceedings of the IEEE/RSJ International Conference on Intelligent Robots and Systems (TROS), 2017.
Jonathan H Tu, I. Scott Griffin, and Clarence W Rowley. Dynamic mode decomposition: An automated procodure for on-the-fly identification of dynamically relevant structuros in complex flows. Journal of Compulational Dynamics,, 122):771 380, 2013.
Elise van der Pol and Max Welling, On the benefits of invariance in neural networks. In International Conference Learning Represenlations, 2019.
Elise van der Pol, Daniel Worrall, Herke van Hoof, Frans Oliehoek, and Max Welling, Mdp homomorphic networks: Group symmetries reinforcement learning. In H. Larochelle, M. Ranzato, R. Hadsell, M. F. Balcan, and H. Lin (eds.), Advances in Neural Information Processing Systems, volume 33, pp. 4190-4210. Curran Associates, Inc., 2020.
M. W. A. Vieweg et al. Data augmentation of 3d medical images. In Proceedings of the Third International Workshop on Medical Computer Vision, 2015.
Dian Wang and Robin Walters. So (2) equivariant reinforcement learning. In International Conference Learning Representations, 2022.
Matthias Weissenbacher, Samarth Sinha, Animesh Garg, and Yoshinobu Kawahara. Koopman q-learning: Offline reinforcement learning via symmetries of dynamics. In International Conference Machine Learning, pp. 23645 23667. PMLR, 2022.
Matthias Weissenbacher, Rishabh Agarwal, and Yoshinobu Kawahara. Sit: Symmetry-invariant transformers for generalisation in reinforcement learning, 2024. URL https://arxiv, orgooog//////.ooooooooooooo htggabbs/240..5002
Matthew O Williams, Ioannis G Kevrekidis, and Clarence W Rowley. Data-driven forecasting of high- dimensional chaotic systems with long short-term memory networks. Proceedings of the Royal Society A: Mathematical, Physical and Engineering Sciences, 20119)))) 2015.
Denis Yarals, Nlya Kostrikov, and Rob Fergus. Data-efficient reinforcement learning with momentum predictive representations. In International Conference on Learning Representations, 2021a.
Denis Yarats, Tly Kostrikov, and Rob Fergus. Tmage augmentation is all you need: Regularizing deep reinforcement learning from pixels. In International Conference Learning Representations, 2021b. URL
Denis Yarats, Rob Fergus, Alessandro Lazaric, and Lerrel Pinto. Mastering visual continuous control: Improved data-augmented reinforcement learning. Tn International Conference Learning Representations, 2022.
URL https:://openrevvie wwtp:/oummmidddd _tt_uumdddd____yeeeeee
Kun Zhang. Interpreting the internal vector representation an analogical neural network. In Proceedings of the 12th annual conference of the Cognétive Science Society, pp. 684-691. Lawrence Erlbaum Associates, 1990.
Kun Zhang and Terrence J Sejnowski. Parallel distributed processing: Explorations in the microstructure of cognition, vol. 1. In Foundations. MIT Press, 1988.
15

---

Published in Transactions on Machine Learning Research (10/2024)
Seen Training
Lava & Water
Training with Crops
(water)
Agent gaina reward by picking up berriea cropa
(80
[80
Agent
Goal
Berries freward = 0.4)
Grops (reward = 0.151
Figure 10: Modified Lava-crossing with additional option for the agent collect rewards by harvesting berries (purple) and crops (green).
A Modified Lavacrossing Environment Add-On
In this section we extend the empirical results on our custom LavaCrossing environment and present evaluations of point (2) below. We modify the environment by adding:
a safe river (water) to the environment which constitutes the only safe pathway through the lava stream, see Figure (5a). Thus the agent must learn to take "swim" in the river which is a color sensitive choice as the lava needs to be avoided.
2. optionally to point (1) we introduce berry-fields as well as crop-fields upon collecting them the agent receives additional rewards.
In evaluating the performance of CiL the environment option (2), which includes both berries and crops, we observe results that are consistent with those obtained in environment option (1).
When data-augmentation is applied, we notice a detrimental effect on the training process. This is in line with our observations from environment option (1), suggesting that data-augmentation might not be beneficial for this specific task.
Furthermore, we observe that the CiL category demonstrates superior performance when compared to the CNN bascline when no data augmentation is applied. This performance advantage is more pronounced in enviromment option (2) than what was observed in environment option (1), depicted in Figure (5b). This indicates that the benefits of using CiL category II become more evident when the complexity of the environment increases, showcasing its robustness and effectiveness in more challenging scenarios.
A.1 Color Generalisation
We evaluate our trained models on a color varied version the environment, see Figure (10); comparing CNN CNN + CiL for categories I,II, and II, respectively. We find improved generalisation of CiL, see Figure (7b).
6The LavaCrossing environment, standard in the MiniGrid toolkit Sehevaller-Boosverrt et al., 2019), requires agents navigate to goal (green square) without falling into lava (orange squares).
16

---

Published in Transactions on Machine Learning Research (10/2024)
cos Φ cos 0
cos & sin 0 sin &b sin @ cos sb sin sin sin db + cos Φb cos vb cos sin s/
cos ø sin @ cos sb + sin @ sin sin eb sin 0 cos db cos di sin v/ cos 0 cos 1p
W =
sin Φb cos 0
(10)
sin 0
B Technical Details
B.1
Implementation of Symmetry Groups
B.1.1
SO(3) group Category II
The group SO(3) represents the group of rotations in three dimensions, and any rotation in this group can be represented using three Fuler angles. Fuler angles three angles introduced by Leonhard Fuler to describe the orientation of body/vecton They describe arbitrary rotations in three dimensions. The three rotations are oflen referred lo roll, pitch, and yaw, especially in the context of aviation and robotics.
Any element of E SO(3), i.e., any rotation matrix, can be expressed in terms of Euler angles. There are multiple sequences axes about which the rotations can take place (like 'XYZ', where the rotation is first about X, then about Y, and lastly about Z), and the choice of sequence can change the specific angles.
For rotation in 3D space, common sequence is the 'Z-Y-X' sequence (yaw, pitch, and roll, respectively):
1. A rotation & (yaw) about the Z-axis.
2. Followed by a rotation 0 (pitch) about the new Y-axis.
3. Followed by a rotation & (roll) about the new X-axis.
Given these angles, the overall rotation matrix p the "Z-Y-X' sequence can be given as p (0(v)),(0)pp(( where:
cos d
sin ø cos &
P=(0) =
sin ø
cos
sin 0
(0)
0
0
sin 0
Cos 0
0
(a(+)
0
cos 1/
sin +b
0
sin sb
cos
When expanded, the combined rotation matrix written as the weight notation W is:7 with the bounds for the Euler angles:
Yaw (Ø): 0 6 < 2x
Pitch (0): - < <
Roll (w): 0 < 1 < 2n
We employ the definition equation 10 to incorporate SO(3) rotation of color-space i.e. category II, where the Euler angles are given by three trainable weights. the bounds are enforces by using the Tangant-Hyperboli function e.g. 0 < Φ < 2n by Φ π + म Tanh(0), where a trainable parameter.
without problems. The most notorious is "gimbal lock," where you lose degree freedom, and it's possible represent all 3D orientations. This why other representations like quaternions sometimes preferred in applications like computer graphics and robotics.
Note: While Euler angles intuitive, they
17

---

Published in Transactions on Machine Learning Research (10/2024)
B.1.2 Symmetric Weight Matrices Category III
For category III the weight matrices W are symmetric. 8 This is accomplished by simple defining a 3 parameter matrix θ and symmetrise ie add the transpose as
W = softmax(()) + softmax(0)))
(11)
We add the softmax function is added to restrict the values of the weights.
Note that symmetric matrices do not form group. The product of two symmetric matrices symmetric if and only if the matrices commutee, meaning
AB = BA.
If matrices A and B are symmetric, then:
(AB)² = AT
Given that A and B are symmetric:
B² B
AT A
So,
(AB)* = BA
For AB to be symmetric, (AB)T must equal AB, so
AB = BA.
B.1.3 Numerical Invariance: torch.svd
In our implementation, have opted for torch. svd over torch. linalg. svd for computing the singular value decomposition (SVD). Our extensive testing indicates that torch.swi offers superior stability in maintaining invariance under orthogonal transformations of the input, denoted as R and L.
However, is important to highlight that numerical algorithms like SVD are prone to challenges associated with sign ambiguity. This means that the columns of the output matrices can undergo multiplication by #l. To tackle this issue and achieve consistent results, we employ torch. abs on the results of equations 7 and 8. This step is imperative to address the sign ambiguity introduced by torch. ensuring that the final output is consistent and unaffected by potential sign changes.
In essence, the adoption of torch.svd and the application of torch, abs contribute to the stability and reliability of our implementation, especially in scenarios involving orthogonal transformations of the input.
B.2 Normalized Rewards in Procgen
See Figure (11) for test curves or Procgen during training.
To compute normalized scores, we leverage a simple yet effective rescaling technique. By using the observed minimum and maximum rewards in the dataset, we ensure that the scores are bounded between and Specifically, given reward value rwd, the normalized score is calculated using the formula:
reward minimal reward
normalized score maximal reward minimal reward
WWe incorporate symmetric matrices in work which is not to be confused with related concept to orthogonal groups the symmetric group. The symmetric group three elements, often denoted Sg, is mathematical structure that encapsulates all possible permutations of three distinct objects. This group has 3! (i.c., distinct clements, each representing unique permutation. When cousidering matrices, natural represcntation of S3 single "1" in every row and column, with all other entries heing "0". Each of these matrices representa permutation of the standard basis vectors p permutations are orthogonal thus is sub-group of 8O(3).
permutation matrices where each matrix has
18

---

Published in Transactions on Machine Learning Research (10/2024)
Awarape Tost Curves Acoua Deatonments
Figure 11: Sample-effiiency of CiL: Procgen testing curves during DrAC training of baselines and CiL, Table We average over all 5 environments and seeds. We show the mean and the mean of the standard error of the individual tasks.
Game
Min
Max
BigFish
1
40
StarPilot
2.5
64
FruitBot
-1.5
32.4
CoinRun
0.5
13
Jumper
1
10
Table 4: Min-max max rewards for Procgen environments used to compute normalized scores in Table (1) and (2).
Here, minimal reward represents the smallest observed reward in the dataset, and maximal reward represents the largest observed reward. This normalization approach ensures that rewards are set relative to their observed range, providing a standardized perspective on their values.
All experiments performed NVIDIA GPU A-100 V-100.
B.3 Hyperparameters
We summarize the hyperparameter choices in Table (5).
C Proofs of Invariance
The goal of this section to prove the invariance of the quantities defined in equations (3). We start by giving some necessary definitions and proving some preliminary results in section (C.1).
C.1 Preliminary steps
Given N > O, we define the following subset of x x B x NNN x NNN
In.N ((S, A,U)SS'S
19

---

Published in Transactions on Machine Learning Research (10/2024)
Table Architecture and syper-paramete choices for CiL on Procgen, DMControl, Minigrid based on (Raileanu et al., 2020), (Hansen & Wang, 2021), and (Jiang et al., 2021), respectively. Channels refer to the category channels. We use the algorithms in code-base without any hyper-parameter changes except reduction of hidden-dim of the actor-critic networks to 64. The patch size follows the convention in Vision Transformers; for 64x64 pixel input, we use 8x8 patches. In Figure 3, the neighborhood sizes are chosen as odd numbers (3, 5) to ensure a central patch with an equal number of neighboring patches on all sides.
Suite
Model
Patch Size
Channels [1.I1,111]
Activation
Hidden Dim.
CNN
3 ResNet
(Taaileam et al., 2020))
[0.0.0]
N/A
64
CiL
8x8 pixel (64 patches)
[1,20,0
sclu
64
CiL Ber
8x8 pixel (64 patches)
1,0,20
sclu
64
(Raileannet al., 2020)
Procgen
CiL!-m
8x8 pixel (64 patches)
[1,10,10]
sclu
64
CiL
8x8 pixel (64 patches)
[1,10.10]
sclu
64
CiLg
8x8 pixel (64 patches)
[1.10.10]
sclu
CiT 1-177
8x8 pixel (64 patches)
[1.10.10]
selu
64
DMGontrol (Honsen& ong, 2021)
CiL!
8x8 pixel (144 patches)
[1,1,0]
selu
64
CIL
7x7 pixel (4 patches)
[100.]]
selu
32
Minigrid (Jianget ol., 2021)
CIT.!
7x7 pixel (4 patches)
0.11.0
sehu
32
CIT.
7x7 pixel (4 patches)
[0.0.1]
sclu
32
is diagonalizable and
(A,U) )) is eigendecomposition sf
Let RER Rax
LE EXN be orthogonal matrices; their orthogonality implics that
(RS') St
(12)
and
S
(13)
Therefore we arrive at the following lemma:
Lemma 1. Let UA, be eigen-decompositionn of St and let R E Ruxn L E NxN be orthogonal matrices. Then
The pair (A, RU) is an eigen-decomposition of (RS))
2. The pair (A, U) is an eigen-decomposition of (SL))
Proof. For the first part, using (12), have
(RS)) R (RU) R(S'stJ,
which is equal to R(UA) (RU)A by the hypothesis. The second part follows immediately by (13) and the hypothesis.
The above lemma allows us to define left action of the group of orthogonal n matrices O, on Fn,N
2n.N On X In,N N In,N (S, U)) + (RS, A, RU)
as well as a right action of the group of orthogonal N x N matrices ON on In.N
VA.N ON x I + In.N (L, (SL--.'L--
20

---

Published in Transactions on Machine Learning Research (10/2024)
We further define the functions
X1 :::.N + R (8, "S,S",A,U) + diag(U"
and
: F - B (5,
Proposition 5.
1. The function X1 invoriant with respect to the group actions pn.N and v n.N.
The function X2 is invoriant with respect to the group action Pa,N.
Proof. 1. For any (S, ,A,,) On In,N, we have
X1 a Ln (R.
(14)
= Xi((SS, RS' A, RU))
diag((RU)) RS (RS) (RU))
diag( R R
diag(U" s'stt
= ((((S,',AAU))
where we used the fact that orthogonal to write (RS) StR and R I.
For any (S,S,,,, E ON In,N, have
X1 (L.N(L,(S,S,AAAU
(15)
(iog(("(SL=))(SLL
diag ("ng(("(("L))
diag diag(U""('S
= (A((S)'
(16)
where we used the orthogonality of L to write (SL-1)f- Lst.
21

---

Published in Transactions on Machine Learning Research (10/2024)
2. For any (R. (S,S",AA) On In,N, have
PN,N Cn (S,
X2
x= RS", RU))
(RS)"
- UT RS'
S²SS
(17)
A singular value decomposition () S is a factorization
$ = UEV*
where U, V are unilary matrices and Σ is reclangular diagonal with non-negative entries on ils diagonal. Given S, such a decomposition need not be unique, however the matrix Σ is unique up to ordering of its diagonal entries (singular values of 8). We will consider singular value decompositions where the diagonal entries of Σ appear descending order. Then, for any r we define U, to be the matrix formed by the first r columns of (also known as the modes).
Lemma 2. Let (U, E, V) be singulor volue decomposition of S and let R E BRxn E NNN be orthogonal
matrices. Then
ER is a singular value decomposition of RS.
2. (U, L-V is singular value decomposition of SL.
Proof. By the hypothesis we have S UEV" Therefore, RS = (RU)EV* and RU is orthogonal. Further, ST. USV = U..(---VVV) and L- is orthogonal.
Next, we introduce the following subset of REXN x BEXN x gaxe x RAXN x NNN
SSVD (.S.S.,.EEEVVV,,,,
is a singular value decomposition of S}. By Lemma 2 able to define left. action of the group of orthogonal matrices On ISVD
On x SVD SVD (S, + (RS, RS', RU.
SVD.
R.N
well right action of the group of orthogonal N N matrices On on ISVD
SVD RNN
ON SVD SVD
R.N
Further, we define the following function
ky SVD - Rr
(S, U,,.U,,VV V) U..
22

---

Published in Transactions on Machine Learning Research (10/2024)
Proposition 6.
The function k, is invariant with respect to the actions D and P.N.
Proof. We have
kr
& D (R, (S,..,,),
(RS, RS" RU,
(RU), (RS)) (RU),
(RS))
(R R
U,SSU,
k,(S,S,U,E.V)
(18)
and
U
S
U,S'S
k,(8,S,UEE.V)
(19)
7
C.2 Completion of proofs
To complete the proofs, usual 2D image is identified with a matrix whose entries correspond to pixel color intensities. In our current setup, we further identify matrices with vectors by stacking together consecutive rows of a matrix, that is
gmxk - Rmn
T11
012
11k
Ø22
311
312
.7/22
321
D22
22k
Tm1
Kink
Fml
Fm
Zmnk
Remark: A rotation and a flip are special cases of pixel permutations, which, under the above identification, correspond to left multiplication by permutation matrix, hence orthogonal matrix.
23

---

Published in Transactions on Machine Learning Research (10/2024)
CiL-Stack
LLL H
S (La) Tet" Waj (LF) Diy (I2) y V(S) jkkm!
(20)
C.2.1 Proof of Proposition (1) and Proposition (4)
We build on the results of section (C.1). Let us repeat proposition (1) below.
Proposition 7 (CiL Symmetry Invariance). Our approach equation 7 is invariant under:
1. right mairia multiplication by orthogonal & equation 5 Wf =1 i.e. the idenlity matric.
2. by pee-parameter family of (special) orthogonal Ls equation 5 if E SO(3). The specific L depend on the details of weights in W, see the appendix section (B.1) for evact relations.
3. by 3 (6) symmetric C2, Ls, equation 5 iff W is symmetric. The specific L depend on the details of weights in W, see section (B.1) for exact relations.
The expression in the parenthesis proposition (1) refer to the case when we do not normalisc by its determinant. In the default case we use V V/det(V).
Proof: The point (1) above follows a consequence of the discussions in section (C.1).
The proof of point (2) and (3) of proposition (1) builds on point (1) thereof, i.e. by adding respective
weights W to preserve specific symmetries and break others. Firstly, for claim (2) note that symmetries L
are preserved if and only if they commute with W E SO(3) ic. [W, L] = WL LW= 0. Where L acts on the input S + L L thus
CIL SL + SL W L
is orthogonal i.e. LLT LT the expression above is invariant if and only if the claim of commutativity between W and L holds.
Note that any real orthogonal 3 x 3 matrix W E SO(3) may be decomposed as W UAU with one real and a pair of two complex conjugate eigen-values all lying on the unit disk (Strang, 2003). Thus the real eigen-values are either {+1, -1} and the complex ones are {a+ib, -i6), a, b e R s.t. a² = 1. Moreover, invertible normal matrices (such as W and L commute they have a shared set of eigen-vectors, ie.. they share the same eigen-decomposition but for different eigen-values i.c. UALU-!
We normalize the network to symmetries with determinant of {+1) by adding to the layer V - V/det(V). Moreover, due to the constraint of lying of the eigenvalues on the unit disk only one independent parameter remains in AL i.c. is not fixed by the constraint of commuting with W. This free parameter in choosing Ag constitutes the one-parameter family of symmetries which is preserved category symmetry type-
The argument for claim (3) of proposition (1) analog to the one for claim (2). With the difference that now W is a symmetric matrix. Any real symmetric matrix can be written as W UAUT for an orthogonal U, which is referred to as the Spectral Theorem (Strang, 2003). Since U is real the eigenvalues of are constraint to {+1, 11} and the additional constrain of determinant one, results in the following four choices for Ar: diag(+1,-1) 1, -1), diag(-1, -1), diag( 11, -1, +1), diag(1, 1, 1). The latter is the identity matrix. If we allow L's with determinant of either {+1, -1} then one obtains eight different choices for L. As we omit plus/mimus the identity matrix this results in 3 (6) choices for L category III
This concludes the proof of proposition (1).
Proof of proposition (4). Let an orthogonal Iransformation act on the input. S + S L and S' - S' L one arrives at equation 20, where we use Einstein summation convention i,j, = 1,..., 3 and T,... , n" = 1, 1...., and where L,, Lg are orthogonal matrices transforming the color and stack dimension, respectively. Any orthogonal may be decomposed in L1, Lz. The proof for L₃ is the same as for proposition (1). As we choose T R N X to be generic the symmetry w.r.t. the orthogonal La is broken general matrix does
24

---

Published in Transactions on Machine Learning Research (10/2024)
not admit orthogonal matrix it commutes with. However, since the weights of T are unconstrained those are free to preserve symmetries in sell-supervised way.
This concludes the proof of proposition (4).
C.2.2 Invariance of equation (3); proof of proposition (2), proposition (3)
By equation (13), equation (3) invariant under the right action of an orthogonal matrix, hence invariant under orthogonal color changes (Figure 4(b)). This implies the statements of proposition (1).
Proposition (3) is corollary of the results of (Weissenbacher al., 2024), by noting that the graph matrix commutes with orthogonal transformations R if and only if R. preserves the graph structure of underlying pixels, up to mirror flips and rotations by 90%.
C.2.3 Rotations, flips, patchwise pixel permutations and multiplications and invertible color changes
Let us discuss the invariance of equations (6). By the above Remark and Propositions 5 and 6, X1 X2 and IC, are invariant under rotations, flips, patchwise pixel permutations and, more generally, patchwise multiplications by an orthogonal matrix R. The functions X1 and ky are also invariant with respect to invertible color changes.
Regional color changes require some additional analysis.
C.2.4 Regional Color Changes
Let
S = (8es, 881 862 Sea Stg, 8t21 S8gg Sty, 88
Since X1 is invariant under matrix permutations (Proposition 5), may assume without loss of generality that the color changes apply to fixed number k of rows in bottom of the matrix S, while the rest of the rows above them are unchanged. In other words, we can define the regional color change operation
regL 12N3N + R2NNN
[ S
S S2
+
S2
where the matrix S E R(n=N)x3N is the upper part of S, the matrix S2 kx33N is the lower part and L is an arbitrary diagonal 3N x 3N matrix, which acts by scaling the RGB channels of the fixed regions in the sequence of images represented by matrix S2 (the notation here refers ordinary matrix multiplication)
Claim 1. The matriz of eigenvalues in equation 7 is in general, not invariant under regL. Proof. Let us choose matrices
[ S = S'
S S
S2
such that
1. S and S, are square orthogonal matrices of the same size.
2. diag(S)ST*) # ding(S,,5,1)
25

---

Published in Transactions on Machine Learning Research (10/2024)
S
0
(I+1P)-1
L. (I
ST 0
(S) (regi(S)))
0
$
(I+ 12) 13. (1+
0
S2
S 1
S (T+ 127)-152
=
1.. (T+ 177
STT, (T+ 137
(22)
Observe that
31
[1
I
reg
182.
S2
T.
therefore
(S') (regi(S))
(21)
0
S1
8
L
0
S2
L
S
1
1
[S
0
5%
[
L
0
82
S
[1]
[1
[ST
5%
L
L
85
Since 2
7
is full rank, we have
L
E - ( 1 E
(3)- L. (L+
so we see that this implies equation 22. From equation 22 we see that the set of eigenvalues of regi(S"):
(regz (S)) not independent of L; indeed, if that were the case then the trace of the above matrix (i.e. the sum of its eigenvalues) would be the same for L. However,
trace(regr.(5") (regn(()))))
(23)
Liag(S), )a
digi S, );s
1 I
1 1 L
which is constant if and only if diag(S) = diag (just take partial derivatives with respect to L.e).
Claim 2. The matrix Kr in equation 6 is in general, not invariant under regL-
Proof. Let us choose S, S' satisfying the assumptions of the proof of Claim 1. Let us also consider the case where is the full rank of S, so that K, K = Stst. Let (U, E, V) be a singular value decomposition of S. Further let (Uz, Ex, Vt.) be singular value decomposition of (S), that is the transformation of S under
26

---

Published in Transactions on Machine Learning Research (10/2024)
regional color changes operation regl defined above. If equation 6 were invariant under regl, then we would have -for all diagonal T-
U'SuU = Ut regi(S") (regı(8))) Un
which would imply that
trace(U*S'SSU)) trace(Už, regn((") (regr((S)) U..
and therefore
trace(5' St) (race(reg)(S") (regı(S))))
for all diagonal L, since trace(AB) = (racce(B)) for any matrices A, B and since U UüULL I.
However saw in the proof of Claim 1, (f)) is not independent of L; contradiction.
27


Reviews:
[{"review_id": "XSqx3kVt1q", "reviewer": "Reviewer_4MkC", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "**Strengths**:\n- The paper is well-motivated and well-organized. \n\n- The authors provide a theoretical proof demonstrating the symmetry invariance of the proposed method.\n\n- The classification of color symmetries and their corresponding visualizations (in Fig.2) are interesting. \n\n- The authors conducted empirical evaluations across several benchmarks.\n\n**Weaknesses:**\n\n- The main concern is that the experimental results do not demonstrate the superiority of the proposed method over the baselines: 1) Regarding the results on the modified Lavacrossing environment (see Fig.5(b)), the shading area of the CNN baseline overlaps with the shading of category II or category III.  2) In Table 3, the scores achieved by CiL are only comparable to the CNN baseline on the two DMC tasks. 3) Furthermore, the authors did not report error bars of the scores obtained on the Procgen tasks (see Fig.9, Table 1 and Table 2). Therefore, it is insufficient to draw the conclusion that CiL enhanced performance and sample efficiency than baselines.\n\n-  The random shift image augmentation presented in DrQ-v2 [1] has achieved promising results in image-based RL tasks. Comparing CiL to DrQ-v2 could improve the paper’s impact. However, experiments comparing CiL to it are missing. \n\n- The three distinct categories perform differently across tasks. Discussions about how to choose the CiL category for the specific task in practice are suggested to be added. \n\n- The authors show the results of CiL with two categories (in Figure 7). I recommend the authors to add implementation details about how to use two categories at a time and discussion about when to use CiL with two categories rather than one category.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper presents a novel neural network based on the Dynamic-mode-decomposition algorithm to improve color generalization in image-based reinforcement learning. Three different categories of color symmetries are introduced in the paper. Each color-symmetry category can be achieved by determining the learnable parameters matrix of the neural network. The paper provides insightful discussions about the unuseful color symmetry. The authors perform empirical evaluations on the modified Lavacrossing environment, and the environments from Procgen, and DeepMind Control suites.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "- Providing additional statistically significant results to support the statement in Subsection 4.1 that CiT improved the performance than the CNN on the modified lava-crossing environment. \n\n- Additional experiments are suggested to compare CiT with the random shift image augmentation by DrQ-v2.\n\n- Presenting error bars for the Procgen results shown in Fig.9, Table 1 and Table 2.\n\n- Explaining the meaning of the error bars in Fig.5 (b). \n\n- It would help to reproduce and understand the work by providing more implementation details, such as those of the baselines and the RL algorithms mentioned in subsections 4.1 and 4.2.\n\n[1] Yarats, Denis, et al. \"Mastering visual continuous control: Improved data-augmented reinforcement learning.\" arXiv preprint arXiv:2107.09645 (2021).", "broader_impact_concerns": "No concerns on the ethical implications.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "Z0FwAH8nFT", "reviewer": "Reviewer_P3Da", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths\n- The paper is well written and understandable and presented.\n- The paper is very well connected to the literature and references all relevant papers.\n- The problem is important, as some degree of color invariance is required in all learned policies, depending on the environment and task, which should improve robustness.\n- The mathematical formulation and explanation of the proposed method is solid, I only have minor remarks for improvement here.\n- The baselines are color jitter and random conv, seems appropriate for the problem of color invariance, and the proposed method outperforms the baselines with ease.\n\nWeaknesses\n- While I understand the formulation for color feature invariance, to me it is not clear why the method is claimed as self-supervised, there is no description of the self-supervised part in the paper, its not clear to me how this method is self-supervised in the common meaning in ML and CV, compared to autoencoders and contrastive learning for example. While the invariance is definitely learned according to the task/environment, maybe the method cannot be called self-supervised.\n- The results are limited as they show the proposed method obtaining better reward, but there is no direct validation about color invariance in the more complex settings (procgen and distracted dm control). Here it would have been nice to provide some visualizations to show the color invariance in multiple settings.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper is about generalization to color variations when learning a reinforcement learning policy. In particular the authors use learned color invariant features by using the singular value decomposition and adding some trainable weights to account for different kinds of invariances, which can be configured by the user.\n\nResults show that in a challenging setting (modified lavacrossing) the proposed method learns better policies that obtain higher accumulated rewards, and this is also confirmed in the procgen and distracted deepmind control benchmarks. Additional experiments in modified lavacrossing confirm that the proposed method learns different kinds of color invariances that are useful in a reinforcement learning setting.\n\nContributions and new knowledge are:\n- The CiL layer for color invariant feature learning in reinforcement learning.\n- A conceptual framework for learnable color invariant features that considers three kind of invariances, that can be set by the user.\n- Experimental results showing that the CiL layer improves performance by learning color invariant features in a three environments using DQN and PPO.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "- Please clarify how this method is self-supervised. To me this is not clear from the paper, i.e. what exact formulation or mechanism makes the model self-supervised, and how did you validate this?\n- I note that experiments show that categories I, II, and III, provide different reward performance, but something that is not clear from the practical/implementation point of view, how should practitioners select which category to use? Should it just be\n- For validation purposes, can you provide samples of the learned invariant features? I would expect different inputs with color variations, and the learned invariant features that are output by CiL in different settings.\n- How is patch size chosen? Of course this would be task specific, and at least the paper should mention how the patch sizes were chosen, so practitioners can tune the patch size for their specific task. I suggest this could be described in the appendix.\n- It would be nice to have clear claims/contributions statements in the papers' introduction, and an audience statement.", "broader_impact_concerns": "The paper presents a generic impact statement, and I think its insufficient to claim no impact. Feature learning definitely has an impact, as the invariances to color can have negative societal impact, just think about skin color, ideally one wants a model to ignore skin color, but depending on the task skin color could be required, and being invariant to it could reduce performance, so definitely there can be ethical implications of this work that could be discussed, as invariances are not neutral. Another example is to simply learn the wrong invariance that works in train/val/test sets, but completely fail in an open set setting  in the real world.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "87flNUL4Hi", "reviewer": "Reviewer_PWdg", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "### Strengths\n* The research question studied in this paper is of importance to the community. Preserving color sensitivity is critical to achieve better generalization in many real-world tasks (e.g. autonomous driving). \n\n### Weaknesses and other questions\n* I don't quite understand how the proposed method can preserve color sensitivity (specifically Eqn. (7)). I would appreciate the authors could provide some intuitive explanations.\n* From Figure 3, it seems that the benefits could also come from patch-wise color augmentation. Have the authors done such ablation experiments?\n* Figure 4: how are \"useful\" and \"less-useful\" defined?\n* Figure 1 (right-bottom subfigure): how to interpret the two axes?\n* For the MiniGrid experiment in Section 4.1, is training and testing on the same environment?", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper introduces a self-supervised approach to improve color generalization in reinforcement learning. Specifically, a Color-invariant Layer (CiL) has been introduced, which incorporates spatial and temporal permutation invariance. Compared to previous data augmentation methods such as random-conv and color-jitter, the proposed approach shows better performance in achieving color generalization while allowing color sensitivity.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "Please refer to the section above.", "broader_impact_concerns": "No", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: Reviewer 4MkC and PWdg were concerned about the practical benefits, since the practical benefits compared to existing method for RL from images that do not explicitly address color invariances have not been demonstrated. Reviewer PWdg is leaning to reject for this reason. However, as all reviewers agree that the claims are substantiated and the findings relevant for some of TMLR's audience, I argue that the paper can be accept.

I request a very minor revision: The citations in the paper are not properly formatted. For the camera-ready submission, please use \citep{} to put the full citation in parentheses, unless the citation is used as a noun in the sentence.
Decision: Accept (poster)