Rollout Total Correlation for Deep Reinforcement Learning
Authors: Bang You
Huaping Liu
Jan Peters
Oleg Arenz
Venue: Accepted by TMLR ()

Abstract: Learning task-relevant representations is crucial for reinforcement learning. Recent approaches aim to learn such representations by improving the temporal consistency in the observed transitions. However, they only consider individual transitions and can fail to achieve long-term consistency. Instead, we argue that capturing aspects of the state that correlate with other states and actions of the trajectory---even more distant in the future---could further help in extracting task-relevant information. Hence, in this paper we investigate how to learn representations by maximizing the rollout total correlation, the correlation among all learned representations and actions within the trajectories produced by the agent. For improving rollout total correlation, we propose to combine two complementary lower bounds based on a generative and a discriminative model, combined with a simple and effective technique of chunk-wise mini-batching. Furthermore, we propose an intrinsic reward based on the learned representation for better exploration. Experimental evaluations on a set of challenging image-based simulated control tasks show that our method achieves better sample efficiency, and robustness to both white noise and natural video backgrounds compared to leading baselines.

Published in Transactions on Machine Learning Research (08/2025)
Rollout Total Correlation for Deep Reinforcement Learning
Bang You
School of Information Engineering Wuhan University of Technology
Huaping Liu"
hplau@tsinghua..duu edu.crr
Deportment of Computer Science and Technology
Tsinghua University
Jan Peters
con.potters@ita-darmmttadd
Intelligent Autonomous Systems Techmische Universität Darmstadt
German Research Center for Al (DFKI)
Hessian Centre for Artificial Intelligence (Hessian.AT) Centre for Cognitive Science (CogSci)
Oleg Arenz
oleg. dag.aarrrndtttdddmsaaddddd de
Intelligent Autonomous Systems Technische Universität Darmstadt
Reviewed OpenReview: https: openreview. net/ forum? id= gdRRALL8.
Abstract
Learning task-relovant representations is crucial for reinforcement learning. Recent ap- proaches aim to learn such representations by improving the temporal consistency in the observed transitions. However, they only consider individual transitions and can fail to achieve long-term consistency. Instead, argue that capturing aspects of the state that correlate with other states and actions of the trajectory- more distant in the future- could further help in extracting task-relevant information. Hence, in this paper we investi- gate how to learn representations by maximizing the rollout total correlation, the correlation among all learned representations and actions within the trajectories produced by the agent. For improving rollout total correlation, we propose to combine two complementary lower bounds basod on a generative and discriminative model, combined with simple and ef- fective technique of chunk-wise mini-batching. Furthermore, we propose an intrinsic reward based the learned representation for better exploration. Experimental evaluations a set of challenging image-bases simulated control tasks show that our method achieves better sample efficiency, and robustness to both white noise and natural video backgrounds compared to leading baselines.
1 Introduction
Reinforcement learning (RL) has achieved impressive success by enabling robots to acquire skills directly from high-dimensional images (Laskin et al.] 2020). A critical ingredient in reinforcement learning from images to learn concise representations of high-dimensional images that relevant to the task at hand. Reinforcement learning agents can subsequently learn policies or Q-functions based obtained task-relevant representations. However, extracting such representations is challenging, since reinforcement learning objec- tives have no direct incentive to filter out redundant or irrelevant information in image observations, such as noise or changes in the background. Furthermore, representation learning in RL deeply intertwined with
CCorrepoondng Author: Huaping

---

Published in Transactions on Machine Learning Research (08/2025)
the problem of exploration (Yarats et al. 2021a; Yuan et al, 20230: learning representations requires diverse observed data of the agent, while effective exploration can only be achieved with expressive representations.
Previous methods alleviate the above challenges by constructing both auxiliary representation learning ob- jectives and reward functions. Particularly, extracting task-relevant representations of observations can be achieved by constructing auxiliary representation learning objectives that capture temporal consistency collected transitions (Pathak al.) effective technique to capture the temporal consistency is to preserve the mutual information related to environment dynamics (Kim et al., 2019a; Tao et al.) 2020 Bai et al. 2021), such as the mutual information between the representation of the next state, and the representation of the current state and the action. However, these prior methods based on mutual-information usually only capture predictive information in individual transitions and may fail to model temporal coherency within sequential observations.
2017 Sekar et al. 2020 Guo et al. 2022) Chiappa et al., 2023 : An
We argue that enforcing temporal consistency throughout longer sequences of states and actions more efficient for learning lask-relevant representations. For example, the energy of an approximately closed system is very powerful feature, precisely due to its temporal consistency, which allows for long-term predictions. By maximizing the correlation within whole trajectories, we aim to extract such important reprosentations from high-dimensional inputs.
Hence, in this paper we propose the rollout total cor- relation (ROTOC) objective to learn task-relevant representations that capture temporal consistency in sequences of states and actions (see Fig. D. Our rollout total correlation objective maximizes the to- tal correlation among latent embeddiings and actions in the trajectories encountered by the agent. The captured total correlation can be naturally inter- preted as the overall information gain that related to the temporal evolution of the enviromment, allow- ing us effectively enforce temporal consistency on soquences of the obtained embeddings. We show that the lotal correlation objective can be lower- bounded by a sum of per-step mutual information objectives that are similar to previous objectives. While this lower bound is more convenient for opti- mization than the total correlation, misses the long-term correlations that we want to capture. Still, we provide empirical evidence that suggests that representation learning based on such per-step objectives can be biased towards capturing long-term correlation, using simple technique that we refer to "chunk-wise mini-batching" ennsuinn that the mini-batches contain several subsequent transitions. Our experiments show that this sampling strategy not only increases the long-term predictability of the learned representations, but also increases the performance of the reinforcement learning agent. Furthermore, we encourage the agent to explore novel states by constructing an auxiliary reward that measures the task-relevant state novelty by a discriminative bound on the predictive information.
Fnooder
Policy
C1
Transition model
Z
a1
a2
Zp
max Clzg) ai ...: (T7)
Figure 1: Our rollout total correlation objective max- imizes the overall interdependencies between latent state embeddings Z and actions in a sequence for learning representations, which are temporally consis- tent throughout the sequence.
The main contributions of our work are summarized as follows. We introduce the rollout total correla- tion objective for representation learning in RL, which aims to learn representations that capture long-term temporal correlations within the trajectory. We simplify our objective into a sun of per-step mutual in- formation, and derive a tractable lower bound of the rollout total correlation, by combining two different lower-bounds of the per-step mutual information. We propose chunk-based mini-batching, a technique for increasing multi-step correlation of the learned embeddings. We propose an intrinsic reward based on our learned representations that quantifies the task-relevant state novelty for guiding exploration. We conduct extensive experiments on a set of challenging image- based continuous control tasks, including standard Mu- joco and Mujoco tasks with distractive backgrounds, where use both, white noise and natural videos. Empirical evaluations show that our method achieves better sample efficiency on high-dimensional Mujoco tasks, while being more robust to observation perturbations than leading basclines.
2

---

Published in Transactions on Machine Learning Research (08/2025)
The remainder of the paper is organized as follows. In Section 2 we discuss prior related work. In Section we present the proposed rollout total correlation objective. We present empirical evaluations in Section In Section G we draw a conclusion and discuss limitations and future work.
2 Related Work
Unsupervised representation learning methods in RL focus on extracting expressive state representations
based on which the agent makes decisions (Yarals et al.) 2021b; Zhu et al.] 2022: You & Liu 2024 Tomar
et al.
2024 Ni et al. (202)). Common techniques proposed to learn representations include using bisimulation
metrics [emertas & Jepson
2022 Zhang ct al., 2020: Castro 2020 reconstructing raw obscrvations Yarats
et al.
2021b Voelcker et al. 2024. and contrastive learning Laskin et al. 2020 You et al. 2022). Particu-
larly, task-relevant representations can be extracted from observations by capturing the temporal consistency
in collected transitions. An effective technique to capture temporal consistency in transitions is to construct
forward or inverse transition models of the environment in the latent space of representations (Sekar et al. Guo et al. 2022: Zhao et al. 2023). Our work is more related to another group, where capturing the temporal consistency is based on preserving the mutual information related to environment dynamics (Loe ett al. 2020 You et al., 2022; Bai et al.) 2021 Lee et al., 2023), such as maximizing the mutual information between the representation of the next state, and the representation of the current state and the action. However, prior methods based on mutual information only capture the temporal consistency in individual transitions involving two successive states given the action. Instead, we aim to capture temporal consistency throughout a trajectory. To this end, we derive our method based on a total correlation objective, propose chunk- mini- bani-batching for maximizing the total correlation, and analyse the relation between long-term prodictability and downstream performance of the RL agent.
2020
Some recent works, DRIBO (Fan & Li, 2022] and RPC (Eysenbach et al. 2021 extend the mutual- information objective to the sequential setting in RL for learning task-relevant representations. Fan et al. Faan & Li 2022) maximize the task-relevant information shared between representations of the obser- vations from different views and meanwhile filter out the lask-irrelevant information not shared by the multi-view observations by using a multi-view information bottleneck loss. While DRIBO also uses sequen- tial data points during training, its objective is not related to the total correlation among embeddings and actions. Eysenbach et al.] (2021) maximize external rewards while minimizing the mutual information be- tween a sequence of states and sequence of its latent representations used by the policy, for limiting the amount of observed information that affects the agent's behavior. RPC introduces an upper bound on the mutual information and use it to construct new algorithm, which uses common objective to optimize the policy and the encoder. While our KL-based lower bound on the total correlation takes a similar form to the upper bound used in RPC, we derived it from a different perspective. Furthermore, we combined it with a lower bound based on a discriminative model, introduced and analyzed the simple yet powerful technique of chunk-wise mini-baachinng proposed an auxiliary reward for better exploration based on our total correlation bound, and optimized the encoder and the policy using separate objectives.
Our work is also related to exploration with intrinsic motivation in RL Yaarat et al. 2021a Huang et al, 2024), since we use an intrinsic reward to incentivize the agent to collect diverse data, which our encoder is based on to induce representations. Some exploration approaches are curiosity-driven, which usually build a latent space of representations that filters out irrelevant information in raw states, and then use the difference between predicted representations and true representations as novelty measurement (Sekar et al. 2020 Guo
ct al. 2022 Hao et al.] )023). The most related approaches are those that apply the information theory principle to distill Lask-relevant information from observations and then capture lask-relevant state novelty
in the obtained representation space for better exploration (Too et al.] 2020 Kim et al. 2019b Bai et al.
2021 Different: from existing (non-sequential) information-theoretii methods [Kim et al. 2019a Tao et al.
2020
Bai et al. 2021) that only capture temporal coherence in individual states or state representations, our total correlation objective aims to enforce temporal consistency on sequential transitions.
The concept of total correlation was recently applicd to reinforcement learning by You et al. (2025), alboit as a novel regularizer, with the aim of biasing the policy toward simpler and more predictable behavior. Instead, our work focuses on representation learning and exploration for reinforcement. learning from images.
3

---

Published in Transactions on Machine Learning Research (08/2025)
Furthermore, the work by You et al 2025) optimizes policy and encoder using loss based two generative models for latent state and action predictions, whereas, for representation learning, we focus on the latent state predictability using a generative and discriminative bound, and use separate objectives for policy and encoder. Furthermore, we introduce chunk-wise nini-batching sampling and a novel intrinsic reward basod on our discriminative total correlation lower bound.
3 Preliminaries and Notation
3.1 Problem Statement and Notation
We consider Markov decision process formulated by the tuple M = (S,A, P,r;7); P,r,)); where S is the state space, A is the action space, P(s++1|8;, ar) is the stochastic dynamic model, r(s, a) is the reward function and y E (0, 1) is the discount factor. At each time step, the agent observes the current state and selects its actions a, based on its stochastic policy т T(a:)s)), and then receives the reward r(st, az). The RL objective is to maximize the expected cumulative rewards Ey Σ where T = (#1, (1, 82, , denotes the agent's trajectory with horizon T. We focus on image-based reinforcement learning, where the state space is in the form of images. We do not explicitly consider partial observability of single image obsorvations, but assume that stacking the most recent images yields a high-dimensional but Markovian state 8, = (O,, 01-1,-'' 01-5+1) where o, is the observed image at time step t.
3.2
Total Correlation
For multiple random variables {z1, (2,,,2,""> , the total correlation (Watanabe) 1960 Studený & Vejnarová, 1998) among these variable C(r1, 22, ... Fna) is defined as the Kullback-Leibbe divergence between the joint distribution of all variables and the product of their marginal distributions p(z1)p(z2).... .......ppp p(æx),
/
p(z) 22. dæ ....
C(11:12) (2n)
p ... 2n) log
p(x,)
X1 XL.X2 X
Total correlation is an effective measure of the total statistical dependence between multiple random variables. However, total correlation suffers from the same problem as mutual information, namely, that directly computing total correlation computationally intractable in general.
4 Rollout Total Correlation
We will now formulate our rollout total correlation objective, and derive tractable lower bound. We also describe how to generate intrinsic rewards, and how method can be integrated with RL algorithm.
4.1 The Rollout Total Correlation Objective
We introduce the rollout. total correlation objective, ROTOC, for capturing temporal correlation within sequence of observations and actions. We first map sequence of high-dimensional states 82,,8,,-.. 1......,0 ST) into compact deterministic representations (2,,c2, "e,....,c, .r) using deterministic convolutional network. We further use parameterized sLochastic neural network that lakes state representations as input and randomly outputs latent samples Z. The goal of ROTOC to find long-term temporally consistent representations Z, given a sequence of states (8), 82, 8T) and actions (a1, .2, ar--))
The proposed objective of ROTOC aims to maximize the total correlation among a sequence of stochastic representations and actions,
max
(1)
which can be defined the Kullback- (KL) divergence between the joint distribution of all embeddings and actions and the product of their marginals,
p(Z1:T 01:r-1
C(21 ai ;ar-1; Ep(z.r.arr-i))
log IT
(2)
[= p(z)) [T= p(a))
4

---

Published in Transactions on Machine Learning Research (08/2025)
4o
Ce
No (c)
(
fo()
0
Mi 9p 12 at)
InfoNCE loss
KL divergence
deterministic encoder (CNN)
projection head (FCN)
prediction (FCN)
Pm
C++1
9m (Zee1 (c+1))
Mm
2c+1
fm()
Om
stochastic encoder
trans ition model
(FCN, Gaussian)
(FCN, Gaussian)
Figure 2: Our network framework. A stack of the most recent image frames is compressed using determin- istic CNN and then fed into stochastic encoder that we use for sampling latent representations of the states. Our lower bounds on the rollout total correlation objective involve learning a latent transition model with a loss given by & KL divergence, and learning non-linear transformations (Zt, ag) and žt+1 that are used by the InfoNCE loss. For the deterministic and stochastic encoder (ø and g) and the nonlinear transformations f, use target networks when processing the state at the next time step 1, and online networks when processing the state at the current time step
The lotal correlation C(21;11:". (T-1; qualifies the total interdependencies among a sequence of stochastic representations and actions, which can be intuitively viewed as the additional amount of in- formation that needs be transferred when encoding every random variable independently instead of using optimal code for transmitting the whole trajectory. Maximizing the total correlation among state embed- dings and actions yields temporally consistent representations that correlate with the agent's 's actions, and, therefore, focus on task-relevant features.
4.2 Algorithm Overview
Unfortunately, maximizing the rollout total correlation directly is computationally intractable, and, hence, derive tractable approximation. We use our approximation learning for learning representations. Furthermore, build an intrinsic reward for guiding exploration.
auxiliary loss during reinforcement
Before introducing the concrete loss functions, the reward function and their derivations, we provide an overview of the proposed network architecture Fig. 2 We use an online deterministic encoder &. with parameters To to extract the state representation e from the state 8, at the current time step. An online stochastic encoder with parameters 0, is used to map the state representation C into mean Ho and standard deviation a of diagonal Gaussian distribution (o(2,)c;) N(H,, diag(o?)), which is used for sampling the stochastic embedding. The sampled stochastic representation 2; and action a, are fed into transition model with paramelers vb, which outpuls the mean Me and standard deviation a of variational Gaussian distribution 4u (*1+1 &t, ax) = NNu,, )iag(o?)) that we introduce for one of our lower bounds. A projection head d, with parameters v and an online prediction head fo with parameters nonlinear transformation to the stochastic representation z, and action a that are used for another lower bound, which relates to InfoNCE, (Oord et al. 2018).. While determining the next. representation C++1 and stochastic representation 2t+1, we use target networks for preventing gradients to flow through the deterministic and stochastic encoders. We also use target prediction head fm without gradient flow when introducing nonlinear transformation of the next representation z++1 for the InfoNCE loss. The parameters of the target networks are not optimized with respect to our objective functions, but follow their online counterparts using an exponential moving average technique (He et al.) 2020). The state representations and their corresponding latent stochastic representations are, hence, created using the models,
are used to introduce a
C₁ (o66 C+1 2m mm); 2t (o(2e(es
Z++1 ~ gm (z |c++:: 0m))
(3)
5

---

Published in Transactions on Machine Learning Research (08/2025)
4.3 Lower-Bounding the Rollout Total Correlation
We will now construct variational lower bound on the total correlation C(21 ;.r-:r)) suitable for optimization. We first simplify our total correlation objective into step-wise losses,
log ((21) I p(2,+1lzt, C(2) -1 ar-1122)) E Bp
T-1
Σ 1 at).
(4)
p(z1) I p(z+11)
Please refer to Appendix A.1 for the derivation. Eq. shows that the total correlation can be lower-bounded Z,, a+) between the current representation 2ạ paired
by the sum of the per-step predictive information [2 21; with the action ae, and the next representation *++1- Although for computationally tractable optimization obtained lower bound in Eq. 4 that does not explicitly capture the long-term correlation, propose to maxmmize the lower bound in multiple consecutive transitions for retaining the long-tern correlation between representations. We will show in Section 5.3.1 that retaining the predictive information in at least two sequential transitions helps to improve long-term predictability and overall performance.
4.4 Two Realizations the Lower Bound
While the lower bound in Eq. is much simpler compared to the exact total correlation objective, it is still not tractable in computation. We will next show two realizations of the lower bound that use different techniques to bound the per-step mutual information. The first objective is based on generative transition model and aims to increase the consistency of the representation, while the second objective is based discriminative model, focusing on the representativeness of the representation.
4.4.1 Lower Bound Based on Generative Model
We obtain a lower bound of the per-step mutual information by introducing a parameterized variational distribution g(z(41 Iz,, az) to approximate the dynamic model in the latent space p(z az):
1(z0+122, 2 - -Bplet DKL p( (((141le442) ((x)=122))a)))
(5)
The detailed derivation available in Appendix A.2 The obtained lower bound given in terms of the expected KL divergence between our stochastic encoder and the transition model qu(zz-1 a.). Maximizing the bound encourages the transition model to approximate the dynamics of the representation Z to predict the next representation well. The resulting latent representations are casily predicted by our transition model. We compute the KL divergence analytically based on the means and standard deviations of two diagonal Gaussian distributions 9m 9m (Z(+1/0(+1) and qo (#141 &t. a,).
4.4.2 Lower Bound Based on a Discriminative Model
Using the KL-based lower bound in Eq. 5 alone could fail to induce meaningful stochastic representations, since for any deterministic embedding, the KT, divergence could always be zero for a constant encoding. To alleviate this problem, we also use the InfoNCE lower bound on the local predictive information term I(z++1; 22, at), which is based discriminative model. While it would be possible to only use the InfoNCE lower bound, we hypothesize that the InfoNCF, lower bound, which favors discriminative representations is less effective in capturing temporally consistent representations, compared to the KL lower bound, which explicitly uses a transition model to predict the future.
Let (ze, a;, 24) denote samples randomly sampled from the joint distribution p(z,, ar R(41) which we refer as positive sample pairs, and let N denote of negative samples 2 drawn from the marginal distribution p(z(+1). Then, the InfoNCE loss (z Zt, ar) is given
(2 ars Z++1)
(2(4+11 22, 21, - p.N
log Exir ENUsm h (Zt, ai 2iti).
(6)

---

Published in Transactions on Machine Learning Research (08/2025)
where the expectation is computed over the joint distribution p(z,, at, zt+1). The score fune- tion Iww 0z, 21+1) transforming feature variables into scalar scores is given by (=(z,, Ot, Ze+1) exp (do(ds do (m(zz, a:)) fm (z4+1) Here the function m(ze, az) concatenates 2, and a and w is a learn- able weight transformation matrix. The projection head (.)+) and the online prediction head fo(-) nonlinearly transform m(z az), whereas the target prediction head fm introduces a nonlinear transformation to Ze+1- By maximizing the inner product between the nonlinear transformations of m( at) and 2++1, the InfoNCE loss with the score function forces our model to learn temporally predictive representations.
4.5 The Loss Function and Optimization
By plugging both the KT. lower bound in Eq. and the InfoNCE: lower bound in Eq. maximize the rollout total correlation (Eq. by minimizing the loss
P into Eq.
we can
Σ log O2 qu (Z141 Zec as)
ha a2 2e+1) at z
arg min min
==
E
hew E log Ex; ENUEL.1 ]
7o.0o:9: Po
(7)
The derived loss contains summation over time steps and enables us to preserve predictive information in time-series observations. We optimize the loss via sampling during training. The learnable parameters of our ROTOC model are simultaneously optimized by minimizing this loss.
4.5.1 Chunk-Wise Mini- Mini-Batching
The loss function (Eq. 7 is optimized using stochastic gradient descent based on mini-batches sampled from the replay buffer of the agent. When sampling individual transitions independently, we would only use single transition for many trajectories, which may prevent us from capturing longer-term correlations. Instead, propose to sample minibatch of sequence chunks (st, ag 8141)(=t/11 with chunk length L from the replay buffer. While our loss (Eq. 7 decomposes additively across time steps, implying that chunk-wise sampling does not alter the optimized loss, it intuitively biases the optimization toward local optima with better long-term consistency by ensuring the mini-batch contains subsequent time steps. This is because gradients from independent single-step transitions can steer the model into local optima exhibiting only local temporal consistency, suboptimal for multi-step prediction. Even if subsequent time steps are eventually sampled in other minibatches, escaping such local optima can be difficult. Instead, chunk-wise mini-batching ensures the aggregated gradient of subsoquent transitions will agrec pushing towards globally consistent representation, effectively canceling out signals that might. otherwise lead to merely locally consistent solutions. In our experiments, we evaluated the gradient variance across the same fixed test trajectories during training, and observed that the gradient variance along the trajectories is indeed smaller throughout optimization and most importantly at the end of it, see Appendix D.5
4.6 Intrinsic Reward for Exploration
By minimizing the loss (Eq. 7, can induce compact latent space of stochastic representations that filters out rastrve information. These representations by themselves can help the agent to more consistently and efficiently explore the environment because filtered information can no longer affect its policy. In this section, we discuss how we can further improve exploration using the task-specific novelty mcasured in the Icarned latent space for constructing an intrinsic reward. Our intrinsic reward function based on curiosity-driven exploration (Guo et al. 2022) where the agent is driven visit novel states with high prediction errors. More precisely, we provide an additional reward the agent based on the InfoNCE loss, which measures how difficult it is to predict the next stochastic representation zi+1 given the current stochastic representation and the action a, The intrinsic reward at the current. time step is given by
(o at 2x+1) Exinh ENUzeN hwl at, 2(+1)
[]
r" 7"( 0:) = ((
log N
(8)
7

---

Published in Transactions on Machine Learning Research (08/2025)
Intuitively, the intrinsic rewards encourage the agent to choose the actions that result in transitions that our model is not able to predict well. We use the discriminative model instead of the generative model g(zx+1/Z,, az) to compute the intrinsic reward, because discriminative models are typically easier to learn. Combing the environment reward r(5;, ar) with the intrinsic reward r" r"(st: az), we obtain the augmented reward r raagg (81; at) = r((,,))) + (xe az) where À is a scaling factor.
4.7 Plugging into Reinforcement Learning Algorithm
We train our representation jointly with soft actor-critic agent (SAC) (Haarnoja et al. 2018). We only feed the first transitions of the sampled sequence chunks into the actor and critic for optimizing the agent. Following prior work (Lee et al., 2020 Bai et al. 2021), the actor takes the deterministic state representations C as input, which empirically outperforms using the stochastic embeddings 2,. We let the gradients of the critic backpropagate through the online deterministic encoder, since the environment reward can provide useful task-relevant information. We stop the gradient of the actor through the embedding, since this can degrade performance by implicitly modifying the Q-function during the actor update (Yarats et al. 2021b) We refer to Algorithm | for the training procedure.
5
Experimental Evaluation
5.1 Experimental Setup
We evaluate the ROTOC on a set of challenging standard Mujoco tasks from the Deepmind control suite (Tassa al. 2018) (Fig. 3 top row). Specifically, we evaluate our method in six mujoco tasks with high-dimensional observations: Ball- in-cup Catch, Cartpole Swingup-sparse, Reacher Easy, Cart- pole Swingup, Walker Walk, and Cheetah Run.
We further carry out our evaluation in two more settings for testing the robustness of our algorithm: the noisy setting and the natural vidco setting. We refer to Mujoco tasks in the noisy setting and the natural video setting as noisy Mujoco tasks and natural Mujoco tasks, respectívely. In the noisy setting, each background image is filled with pixel-wise Gaus- sian white noise (see Fig. 3 middle row) that is regarded as strong state distractor for reinforcement learning tasks (Zhou ct al. 2023). Furthermore, in order to simulate robots in nat- ural environments with complex observations, in the nal ural video setting (see Fig. 3 bottom row) the background of the Mujoco tasks is replaced by natural videos (Zhang et al. 2020) sampled from the Kinetics dataset (Kay et al. 2017).
Cartpole
(b) Rall-in-cup
(c) Rescher
(d) Walker
(c) Cheerah
Casch
easy
Walk
Run
Figure 3: Image-based control tasks used
in our experiments. The images show task observations of standard Mujoco tasks (top row), noisy Mujoco tasks (middle) and natu- ral Mujoco tasks (bottom).
The backgrounds of natural Mujoco tasks are continuously changing during training and evaluation, which introduces realistic and strong perturbations to observation images.
We compare ROTOC with the following leading baselines in our experiments: 1) Proto-RL Yaratt et al, 2021a), which extracts prototypical representations from observations and builds an entropy- intrinsic reward for exploration, 2) DB (Bai al. 2021), which learns dynamic-relevant compressed representations by using an information bottleneck objective and constructs intrinsic reward function based on information gain to boost exploration, 3) Plan2Exppore (Sekar et al.] 2020), which learns task-relevant representations by modeling dynamics and utilizes the disagrecment in the predicted next state representations as intrinsic rewards, 4) RPC (Eysenbach et al.) 2021),, which learns predictable representations from observations and constructs an intrinsic reward function based on the KL divergence, 5) SAC (Haarnoja et al. 2018), which learns policies by maximizing the policy entropy and rewards. We use the same codebase of the SAC algorithm for DB, Proto-RL, RPC, and SAC to ensure a fair comparison other model-free methods, while following the original implementation of the model-based method Plan2explore. The implementation details
8

---

Published in Transactions on Machine Learning Research (08/2025)
Cartpole Swingup
Ball-in-cup Catch
Reacher Easy
1000
1000
800
Lo000de reward
Epis reward
Episode reward
600
600
600
400
200
0-
0-
Cheetah Run
Walker Walk
Cartpole Swingup sparse
800
Elssde reward
Episode reward
Episode reward
500
400
0-
0.5
0.6 0.5
0.0
0.25
0.10
0.55
0.20
0.25
0.30
0.5
0.1
02
0.3
Envinronment stepsix109)
Envinronment stepsix109)
Envinronment stepsix109)
ROTOC
DB
Proto-RL
Plan2explore
RPC
SAC
Figure 4: Learning curves our method and baselines. The plot shows the average reward over 10 seeds and 90% confidence interval. ROTOC outperforms the basclines on the majority of the standard tasks from the Deepmind control suite in terms of sample efficiency.
of our method and basclines are available in Appendix B
We select the coefficienss a and x, and the chunk length I. by performing hyperparameter tuning on the standard Cartpole Swingup task and subsequently fix them for all other tasks. We refer to Appendix for more details on determining hyperparameters.
5.2 Performance Evaluation
In our first set of experiments we evaluate the final performance and sample efficiency of the different mots in Mujoco tasks in the standard, noisy and natural setting.
5.2.1 Standard Setting
We first study whether ROTOC can succood in challenging Mujoco benchmark tasks when using unperturbed images. In Fig. 4 we compare our algorithm with DB, Proto-RL, Plan2explore, RPC, and SAC on six standard Mujoco tasks from the Deepmind control suite. ROTOC achieves better sample efficiency compared to all baselines on five out of six tasks. The results indicate that the learned representations and the intrinsic reward learned by ROTOC significantly help to learn good policies on challenging image- based tasks.
5.2.2 Robustness to Noisy Observations
To study the robustness white noise, we corrupt the image observations the Mujoco tasks with Gaussian noise (Fig. middle row), introducing random task-irrelevant patterns the raw observations. In the top half of Table I we compare the performance at fixed number of environment interactions on the noisy Mujoco tasks. ROTOC achieves competitive asymptotic performance at 500K environment steps compared to the basclines across all tasks. Scores in Table | are average rewards and standard errors over 10 seeds. Learning curves and an additional plot and table to compare performance drops caused by the presence of white noise are shown in Appendix [ The performance of our method remains consistent when handling noisy Mujoco tasks, whereas the basclines show a clear decline in performance. For example, ROTOC achieves similar performance in standard and noisy Ball-in-cup Catch and Cartpole Swingup-sparse tasks, while the performance of strong baselines DB, RPC, and Proto-RL. clearly decreases due to the noise.
9

---

Published in Transactions on Machine Learning Research (08/2025)
Table 1: Scores (mean and standard error for 10 seeds) achieved by our method and baselines at 500k
environment steps on six noisy Mujoco tasks and the natural Mujoco tasks. Our method achieves better or least comparable performance baselines.
500K step scores
ROTOC(Ours)
DB
Proto-Ri
Plan2Explore
RPC
Cartpole Swingup
869
851 + 7
8464 5
773:128
814+13
Ball-in-cup Catch
9644
782 = 155
9224 16
892:4433
928-220
Reacher Easy
947 20
939 + 19
869 63
851 132
980 29
Noisy Mujoco tasks
Cheetah Run
332::
265 20
268 13
2514 22
282+188
Walker Walk
935::
880 30
9131 15
752436
893114
Cartpole Swingup-sparse
768 14
450 = 133
571 95
680 51
72+ 91
Cartpole Swingup
8441
803 16
7484 28
1434288
8124:00
Ball-in-cup Catch
947 12
916 + 17
7274 53
133+136
838472
Reacher Easy
399
327 65
257 39
362:: 192
327 73
Natural Mujoco tasks
Cheetah Run
308 16
294 16
296 8
1934 32
279
Walker Walk
8864 15
841 1 31
8824 20
557454
714--66
Cartpole Swingup-sparse
118:1
9
654 27
0 0
5 4
Table 2: Performance of our model ROTOC for 4 different chunk lengths L on the standard Cartpole Swingup-sparse and Ball-in-cup Catch. The bold font indicates the highest average rewards among all methods. Overall, using larger chunk lengths achieves higher average rewards.
Scores
L=1
L-2
L-3
L=55
Cartpole Swingup-sparse Ball-in-cup Catch
601 + 96
793 + 10
804 + 10
796 + 12
953 + 4
957
965 + 963 + 4
Table 3: Ablation of ROTOC over chunk length L and batch size BS. The table shows the mean and standard error for 10 seeds. When using the same amount of training samples, e,g. 512, 768, and 1280, using sequential transitions achioves better performance than using individual transitions (L -L=1 1).
512 samples
768 samples
1280 samples
Configurations
I==1, BS=512
T-2, BS=256
T==1, BS=768
T.=3, BS=256
I==1, BS=1280
T.=5, BS=256
Scores
707: 25
793. 10
696 18
8041 10
712.t 24
796.1 12
5.2.3 Robustness to Natural Background
We further investigate the robustness of our method on the natural Mujoco tasks (Fig, bottom row), where the natural backgrounds of the observations are constantly changing. The performance at 500K environment steps is shown in the bottom half of Table I and additional plots of the results are provided in Appendix C Our method achieves better or at least comparable performance across all tasks compared to the baselines. Moreover, the performance of baselines (e.g. Proto-RI. and Plan2Fxplore) declines significantly more than our method on the majority of natural Mujoco tasks. The experimental results in the natural video setting show that ROTOC is more robust to task-irrelevant background than baselines.
5.3 Ablation Studies
We performed several evaluations to investigate the effect of chunk wise mini- tini-batching, the intrinsic reward, and the different lower bounds. We also visualize the learned embeddings to investigate which features are encoded. We perform 5 independent runs for all ablation studies in this section, unless otherwise specified.
5.3.1 Chunk-wise Mini-batching
We performed ablations inspect whether preserving predictive information soquential transitions im- proves the performance compared using individual transitions. In all previous experiments we used the fixed chunk length I. =2 2. In the following experiments first evaluated the performance of ROTOC when
10

---

Published in Transactions on Machine Learning Research (08/2025)
Table 4: Future predictability of the learned representations with different chunk length L and batch size BS. For one-step predictions (n = 1), the representations of an encoder that was trained with chunk length L 1 can be better predicted in terms of log-likelihood. However, higher chunk lengths yield representations that are better suited for multi-step predictions (n 2 2).
loelliiel
n=1
11=2
n=3
n=5
T, 1.=1,BS-220 1RRS=226 256
-78.9 + 0.6
-98.3 + 0.3
-106.1
0.2
-114.2
0.4
L-1,BBS512
-78.8 + 0.2
-97.3 + 0.6
-105.6 1 0.1
-115.6 0.5
L--2BSS228
-80.9 + 0.4
-91.9 + 0.4
-99.6 + -0.1
-109.3 0.4
L=3,BS=256
-82.5 = 0.3
-92.5
= 0.4
-98.9
+ 0.1
-107.1 = 0.1
L-5,BS-256
-118.3 + 0.4
-125.4
+ 0.3
-131.4 + 0.0
-139.8 + 0.8
using the same batch size, but different chunk lengths. When the chunk length L is set to 1, ROTOC only uses individual transitions. Using at least two sequential transitions outperforms ROTOC with chunk length L (Soe Table Z). We also obscrve that ROTOC using chunk length of 5 does not further improve performance. We hypothesize that a large chunk length either decreases the variability within the batches (if we compensate for the chunk length by decreasing batch size) leads to too large batches (if we do not compensate), which may result in convergence to worse local optima (Keskar et al. 2017 LeCun et al. 2002) Furthermore, we investigated the performance of ROTOC when using the same amount of data samples, but different chunk lengths. We ran experiments with L = using doubled, tripled, and 5-times larger batch sizes, and compared them to L = 2, L = and L = with a fixed batch-size on the Cartpole Swingup-sparse task. The only difference to our experiment with larger chunk lengths is that the latter guarantees that the minibatch contains consecutive time steps. We can stilll observe a significant performance benefit of batching (see Table 3), suggesting that the different composition of the driving factor for the improved performance of ROTOC. These results are in line with our hypothesis that chunk-wise mini banihbatchin can increase total correlation and thereby increases performance.
chunk-wiss mini-batching
is the
We further analyze why using sequential transitions im- proves the performance. To test our hypothesis that us- ing larger chunk lengths increases correlation between more distant time steps, first learn different encoders using different chunk lengths, and then evaluate the predictability of the learned representations in terms of the log- lieeloood loss of a neural network that we train on the respective data sets. The noural network takes the current stochastic or deterministic represen- tations at time step as inputs and predicts their cor- responding embeddings in the future time step + with prediction step n. Please note, that the differ- ent data sets only differ due to the different encoders, as we use the same underlying data set of states S, which is taken from the replay buffer of the agent that was trained with chunk length L = 1. Table ! and additional results Appendix D.I compare the future predictability of the learned stochastic and determinis- tic representations with different chunk lengths the natural Cartpole Swingup-sparse task. Experimental results show that the stochastic and deterministic em- beddings learned by using chunk length L = 2,3 are better suited for multi-step prediction, but worse for single-step prediction. We attribute the performance gain achievod by using the larger chunk length improved long-term predictability of the learned embeddings by preserving predictive information in mul- tiple consecutive time steps. With larger chunk length L 5, the future predictability of the learned
0.20
in [355 ,5'' and in [-0.25, 0.25) in [175 1885 and [-0.25, 0.25)
Intrinsic reward
0.15
0.10
0.05
0.00
0.2
0.4
0.6
0.8
1.0
Envinronment steps(x109)
Figure 5: Our intrinsic reward function allocates higher intrinsic rewards to the novel states with po- sition in the range of [-0.25, 0.25] and pole angle in the range of [355", compared to the frequently visited states with position in the same range and pole angle in the range of [175', 118") on the Cart- pole Swingup-spane sparse task with natural background.
11

---

Published in Transactions on Machine Learning Research (08/2025)
Cartpole Swingup-sparse
Cartpole Swingup-sparse (White Naise)
Cartpole Swingup sparse (Natural Video)
800
Epsode reward
Episode reward
600
Epsode reward
600
Ro0
2200
200
6.0
0.2
03
0.1
02
03
02
0.6
1.0
Environment stepstx109)
Environment steps(x10f)
Environment steps(x10"))
ROTOC
No-InfoNCE ROTOC
No-KL ROTOC
No-Reward ROTOC
Vanilla SAC
Figure 6: Performance of ablations on Cartpolc Swingup-sparse task in the standard sctting (left), noisy setting (middle) and natural setting (right).
embeddings degrades compared using individual transitions. As we discussed above, we hypothesize that too large chunk lengths lead to too small variance in the stochastic gradient estimates, which may hurt the performance
5.3.2 Intrinsic Reward
fixed data We also study whether the intrinsic reward function effectively guides the agent to explore novel states the presence of natural backgrounds. Specifically, on the natural Cartpole Swingup-sparse task, we store snapshots of our intrinsic reward function every 100,000 steps and evaluate the rewards set. We split the data set into two subsets, where one of the subsets only contains particularly relevant transitions and the other subset contains all remaining transitions. We consider states with a pole angle in the range of [355", 5"], which corresponds to an upright position, as particularly relevant (see more details in Appendix D.4.1) ). Fig. 5 shows the mean and standard deviation of the intrinsic for the two subsets, cach for the same 10 dilferent snapshots of the intrinsic reward function. We also show the evolution of intrinsic rewards along near-optimal trajectories sampled from the fixed data set Appendix D...2 The results show that our intrinsic reward function allocates higher reward values to the novel states than the frequently visited states, indicating that our intrinsic reward can effectively guide exploration in the presence of natural backgrounds.
5.3.3 Individual Contributions of Each Component
We analyze the effects of the individual contributions on three ablations: No-Reward ROTOC lacks the intrinsic reward signal, No-KL, ROTOC lacks the first terms in Eq. and No-InfoNCE ROTOC lacks the second terms in Eq. and does not use intrinsic rewards. We also evaluate vanilla SAC, which uses the same convolutional encoder and SAC architecture ROTOC, without any additional representation learning loss, intrinsic rewards, or chunk- wise mini-bre hami-batching. Fig. G shows the results (mean and 90% confidence interval over seeds) on the standard Cartpole Swingup-sparse task with raw images, white noise, and natural video background. Appendix D.2 shows the results of all three sparse- reward tasks. ROTOC achieves better or comparable performance than its ablations on all tasks. Appendix D.3, we also show the effect of the KL lower bound on the long-term predictability akin to the experiments for Table 4 The results suggest that our lower bound based generative predictive model further increases total correlation.
5.3.4 Visualization
We visualize our learned representations on the natural Mujoco tasks using the Stochastic Neighbor Embodding technique (t-SNE) (Maaten & Hinton 2008) for projecting the representations into 2D plots. To avoid overlapping, we quantize I-SNE points into 2D grid of size 20 15 using RasterPairy (Kingemannn)
The variance across different. stochastic gradient estimates is not confused with the variance of the per-time-step contributions within single stachastic gradient estimate.
12

---

Published in Transactions on Machine Learning Research (08/2025)
Figure 7: We visualize the representations learned by method the Walker Walk task with natural background using t-SNE on 20 x 15 grid. The two images marked in red squares that have different robot configurations and similar backgrounds are far away from each other, while the image pair marked by green squares that have similar robot configurations but different backgrounds, are close to each other.
2015). Fig. 7 shows the visualization of the representations learned by ROTOC on the Natural Walker Walk task. Visualizations for the deterministic embeddings and the baselines shown in Appendix F The latent space of the learned compressed representation organizes the variation in robot configurations even with natural backgrounds. Specifically, images with similar robot configurations and clearly different backgrounds appear close to each other (e.g. the image pair marked in green), while observations with different robot, configurations but similar backgrounds are far away from each other (e-g. the two images marked red), indicating that the learned representations meaningfully capture task-relevant information related to robot configurations, while filtering out task-irrclevant perturbations.
6 Conclusion and Discussion
We presented an inforoaaion-Cheereiiic representation learning method based on rollout total correlation for image-based RL, ROTOC. By preserving the overall interdependencies among sequence of observations and actions, ROTOC can learn temporally consistent representations that focus task-rolevant information, increasing the robustness to noise. For achieving tractable optimization, derived lower-bound the rollout total correlation, which is given by the sum of per-step mutual information. Pursuing our original objective of maximizing rollout total correlation, we proposed chunk-wise mini-batching and adding a more explicit lower bound based on a generative model on top the discriminative InfoNCE bound. Our experiments show that ROTOC indeed outperforms leading methods in the challenging setting of RT. from noisy images. More importantly, our experiments suggest connection between chunk-based mini- batching, long-term predictability and downstream performance, supporting our hypotheses that chunk-wise mini-batching can increase total correlation, and that maximizing total correlation can result in better representatioss.
However, by approximating rollout lotal correlation by a sum of per-step mutual information, our lower bounds may lase some incentive to capture long-term correlations. While the per-step objectives are easier intograte existing RL frameworks and the learned transition model could be used for modcl-basod RI. in future work, it would also be interesting explore objectives that better capture multi-step correlations based multi-step prediction models.
13

---

Published in Transactions on Machine Learning Research (08/2025)
Acknowledgments
This work supported by the National Key Research and Development Program of China Project number 2024YFB4708000, the German Research Foundation (DFG) Project number PE 2315/18- 18-11, and the German Federal Ministry of Research, Technology and Space (BMFTR) Project number 01IS23057B. This project has been supported by a hardware donation by NVIDIA through the Academic Grant Program.
References
Chenjia Bai, Lingxiao Wang. Lei Han, Animesh Garg, Jianye Hao, Peng Liu, and Zhaoran Wang. Dynamic bottleneck for robust self-supervised exploration. Advonces in Neural Information Processing Systems, 34:17007-17020, 2021.
Pablo Samuel Castro. Scalable methods for computing state similarity in deterministic markov decision processes, In Proceedings of the AAAI Conference on Artificial Intelligence, 2020.
Alberto Silvio Chiappa, Alessandro Marin Vargas, Ann Huang, and Alexander Mathis. Latent exploration for reinforcement learning. Advances in Neural Information Processing Systems, 36, 2023.
Ben Eysenbach, Russ R Salakhutdinov, and Sergey Levine. Robust predictable control. Advances in Neural Information Processing Systems, 34, 2021.
Jiameng Fan and Wenchao Li. Dribo: Robust deep reinforcement learning via multi-view information boutlnnek. International Conference on Machine Learning, pp. 6074 6102. PMLR, 2022.
Zhaohan Guo, Shantanu Thakoor, Miruna Pislar, Bernardo Avila Pires, Florent Altché, Corentin Tallec, Alaa Saade, Daniele Calandriello, Jean-Bastien Grill, Yunhao Tang, et al. Byol-explore: Exploration by bootstrapped prediction. Advances in Neural Information Processing Systems, 35:31855 31870, 2022.
Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with stochastic actor. In International Conference Machine Learning, pp. 1871. PMLR, 2018
Danijar Hafner, Timothy Lillicrap, Jimmy Ba, and Mohammad Norouzi. Dream to control: Learning behaviors by latent imagination. In International Conference on Learning Representations, 2019a.
Danijar Hafner, Timothy Lillicrap, Tan Fischer, Ruben Villegas, David Ha, Honglak Tee, and James Davidson. Learning latent dynamics for planning from pixels. In International Conference on Machine Learning, pp. 2555 2565. PMLR, 2019b.
Jianye Hao, Tianpei Yang, Hongyao Tang, Chenjia Bai, Jinyi Liu, Zhaopeng Meng, Peng Liu, and Zhen Wang. Exploration in deep reinforcement learning: From single-agent to multiagent domain. IEEE Transactions on Neural Networks and Learning Systems, 2023.
Kaiming He, Haoqi Fan, Yuxin Wu, Saining Xie, and Ross Girshick. Momentum contrast for unsupervised visual representation lcarning. In Proceedings of the IEEE/CVF conference on computer vision and pallern recognilion, pp. 9729 9738, 2020.
Kaichen Huang, Shenghua Wan, Minghao Shao, Hai-Hang Sun, Le Gan, Shuai Feng, and De-Chuan Zhan. Leveraging separated world model for exploration in visually distracted environments. Advances in Neural Information Processing Systems, 37:82350 82374, 2024.
Will Kay, Joao Carreira, Karen Simonyan, Brian Zhang, Chloe Hillier, Sudheendra Vijayanarasimhan, Fabio Viola, Tim Green, Trevor Back, Paul Natsev, et al. The kinctics human action vidco datasct. arXiv e-prinls, pp. arXiv 1705, 2017.
Mete Kemertas and Allan Douglas Jepson. Approximate policy iteration with bisimulation metrics. Trans- actions Machine Learning Research, 2022.
14

---

Published in Transactions on Machine Learning Research (08/2025)
Nitish Shirish Keskar, Jorge Nocedal, Ping Tak Peter Tang, Dheevatsa Mudigere, and Mikhail Smelyanskiy. On large-batch training for deep learning: Generalization gap and sharp minima. In 5th International Conference on Learning Representations, ICLR 2017, 2017.
Hyoungseok Kim, Jaekyeom Kim, Yeonwoo Jeong, Sergey Levine, and Hyun Oh Song. Emi: Exploration with mutual information. In International Conference on Machine Learning, pp. 3360-3369. PMLR, 2019a.
Youngjin Kim, Wontac Nam, Hyunwoo Kim, Ji-Hoon Kim, and Gunhee Kim. Curiosity-botllenckk Ex- ploration by distilling task-specific novelty. In International Conference on Machine Learning, pp. 3379 3388. PMLR, 2019b.
Mario Klingemann. Rasterfairy, 2015. URL httppp///ggaaa...co/////aaaoaoos
Michael Laskin, Aravind Srinivas, and Picter Abbeel. Curl: Contrastive unsupervised reprosentations for reinforcement learning. In International Conference on Machine Learning, pp. 5639 5650. PMLR, 2020.
Yann LeCun, Léon Bottou, Genevieve B Orr, and Klaus-Robert Müller. Efficient backprop. In Neural networks: Tricks of the trade, pp. 9-50. Springer, 2002.
Kuang-Huei Lee, lan Fischer, Authony Liu, Yijie Guo, Honglak Lee, John Canny, and Sergio Guadarrama. Predictive information accelerates learning in rl. Advances in Newral Information Processing Systems, 33:1 1511880011901, 2020.
Young Jae Lee, Jaehoon Kim, Mingu Kwak, Young Joon Park, and Seoung Bun Kim. Stacore: Spatio- temporal and action-based contrastive representations for reinforcement learning in atari. Newral Net- works, 160:1-11, 2023.
Laurens van der Maaten and Geoffrey Hinton. Visualizing data using t-sne. Journal of machine learning research, 9(N77):2579 2605, 2008.
Tianwei Ni, Benjamin Eysenbach, Erfan Seyedsalehi, Michel Ma, Clement Gehring, Aditya Mahajan, and Pierre Lue Bacon. Bridging state and history representations: Understanding self-prrdietive rl. In 12th International Conference on Learning Represenlations, ICLR 2024, 2024.
Aaron den Oord, Yazhe Li, and Oriol Vinyals. Representation learning with contrastive predictive coding. arX1w:1807.03748, 2018.
Deepak Pathak, Pulkit Agrawal, Alexei A Efros, and Trevor Darrell. Curiosity-driven exploration by self- supervised prediction. In International Conference Machine Learning, pp.
PMLR, 2017.
Ramanan Sekar, Olch Rybkin, Kostas Danillidis, Picter Abbeel, Danijar Hafner, and Deepak Pathak. Plan- ning to explore via sell-supervised world models. In International Conference on Machine Learning, pp- 8583 8592. PMLR, 2020.
Milan Studený and Jirina Vejnarová. The multinffrmation function as a tool for measuring stochastic dependence. Learning in graphical models, pp. 261-297, 1998.
Ruo Yu Tao, Vincent Lavee,, and Joelle Pineau. Novelty search in representational space for sample efficient exploration. Advances in Neural Information Processing Systems, 33, 2020.
Yuval Tassa, Yotam Doron, Alistair Muldal, Tom Frez, Yazhe Li, Diego de Las Casas, David Budden, Abbas Abdolmaleki, Josh Merel, Andrew Lefrancq, et al. Deepmind control suite. arXio preprint wPXXw:1801.00690 2018.
Manan Tomar, Riashat Islam, Matthew Taylor, Sergey Levine, and Philip Bachman. Ignorance bliss: Robust control via information gating. Advances in Newral Information Processing Systems, 36, 2024.
Claas Voelcker, Tyler Kastner, Igor Gilitschenski, and Amir-massoud Farahmand. When does self- prediction help? understanding auxiliary tasks in reinforcement learning. In Reinforcement Learning Conference, 2024.
15

---

Published in Transactions on Machine Learning Research (08/2025)
Satosi Watanabe. Information theoretical analysis of multivariate correlation. IBM Journal of research and development, ((1):66 82, 1960.
Denis Yarats, Rob Fergus, Alessandro Lazaric, and Lerrel Pinto. Reinforcement learning with prototypical representations. In International Conference on Machine Learning, pp- 11920-11931. PMLR, 2021a.
Denis Yarats, Amy Zhang, Πya Kostrikov, Brandon Amos, Joelle Pineau, and Rob Fergus. Improving sample efficiency model-free reinforcement learning from images. In Proceedings of the AAA! conference on Artificial Intelligence, number 12, pp. 10674-10681, 2021b.
Bang You and Huaping Liu. Multimodal information bottleneck for deep reinforcement learning with multiple sensors. Newral Nehworks, 176: 106347, 2024.
Bang You, Oleg Arenz, Youping Chen, and Jan Peters. Integrating contrastive learning with dynamic models for reinforcement learning from images. Neurocompulingg, 2022.
Bang You, Puze Liu, Huaping Liu, Jan Peters, and Oleg Arenz. Maximum total correlation reinforcement learning. In Forty-second International Conference on Machine Learning, 2025.
Mingqi Yuan, Bo Li, Xin Jin, and Wenjun Zeng. Automatic intrinsic reward shaping for exploration in deep reinforcement learning. International Conference on Machine Learning. PMLR, 2023.
Amy Zhang, Rowan Thomas McAllister, Roberto Calandra, Yarin Gal, and Sergey Levine. Learning invari- ant representations for reinforcement lcarning without roconstruction. In International Conference on Learning Representations, 2020.
Yi Zhao, Wenshuai Zhao, Rinu Boney, Juho Kannala, and Joni Pajarinen. Simplified temporal consistency reinforcement learning. In International Conference on Machine Learning, pp. 42227-4224 PMLR, 2023.
Ziyuan Zhou, Guanjun Liu, and Mengehu Zhou. robust mean-field actor-critic reinforcement learning against adversarial perturbations on agent states. IEEE Transactions on Neural Networks and Learning Systems, 2023.
Jinhua Zhu, Yingee Xia, Lijun Wu, Jiajun Deng, Wengang Zhou, Tao Qin, Tie Yan Liu, and Houqiang Li. Masked contrastive representation learning for reinforcement learning. IEEE Transactions on Pattern Analysis and Machine Intelligence, 45(3):3421-344) 34333, 2022.
16

---

Published in Transactions on Machine Learning Research (08/2025)
A Bound Derivation
A.1 Derivation on Eq.
In this section, we show how to obtain the lower bound in Eq. We first introduce the entropy H (arrT-1) the definition of the total correlation in Eq. 2
log p(Z1:T1 a1:T-1)
C(21:11 ;:r-1:=)) By(eur 011-11)
+ H(a1T-11) - H(a1T--)
[T p (2))
[[ p(a))
T-1
Bp(aur 1)
log
p(Z ,(21,T,01TT1)
DRL (f(aIT-I) |I II pla.)
(9)
p(a11T-1) 1= p(z))
t=1
p(z1.T| (1:T-1)
Bp(zrr
log
II p(z))
where the inequality follows from the non-negativity of the KL divergence. We then introduce variational distribution
T-1
((11:( = p(z)) II p(z(+1lz), at);
(10)
which corresponds to Markovian approximation of the conditional distribution p(zur a1.-11). We intro- duce Eq. 10 into Eq. to obtain step-wise losses,
C(2naa1) aa--i:=)) 2 Bp(anr,anr--)) log q 21:7 (1:--1)
IT: p(z))
(11)
log (=z1) ) [T-1 p a.)
T T-1
Ep(tur,
( [ p(z(41)
A.2 Derivation Eq. 5
For obtaining lower bound of the per-step mutual information, start by replacing the entropy H (2 with the conditional entropy H ((x+1104441,
I(z141;21, az) = H(z+11) (((z)))H(zz+122) a.) 2 H H(2411444) H(244122,,0
(ep..4). log p( | ax) p(z (c+11)
(12)
where the inequality holds because additional information can not increase entropy, ( 2 ((
Then, we use variational distribution q(z+11/21,0 a.) to approximate the conditional distribution p(z411/z, at), and then we obtain KL-basod lower bound Eq. 5
log g( 1z az)
+ DKL p(zet1lzz, ar) ll ((2441lZ1, ai)
[ az) >
p (c41)
p(21t1/c44))
Epp DKL (p(ees)lett)) ((2041122,,))) a.))
log q( 22, az)
(13)
B Implementation Details
In this Section, we explain the implementation details for ROTOC in the Deepmind Control suite setting.
17

---

Published in Transactions on Machine Learning Research (08/2025)
B.1 Pixels Preprocessing
Following (Yarats et al. 2021b), we obtain an individual state by stacking 3 consecutive frames, where each frame is an RGB rendering image with size 84 x 84 x 3 from the Oth camera. The pixel values range from [0, 255] and we divide each pixel value by 255 to scale it down to [0, 1) range. We follow Proto-RL (Yarats et al. 2021a) by performing data augmentation by randomly shifting the image by [-4,4], before we feed states into the encoders.
B.2 Network Architecture
The network architecture for ROTOC consists of five modules: deterministic encoder, stochastic encoder, transition model, projection head and prediction head. All five modules are built based on common prac- tice (Yarats et al.) 2021a; You et al.) 2022 Eysenbach et al. 2021) without specific hyperparameter tuning. We also did not tune the network architectures for the baseline methods, but either used the architectures from the original implementation or the same architecture that we used for testing ROTOC. Deviating from the original network architectures was in some cases necessary, particularly when the baseline method was not evaluated for RL from images. In the following, we list and motivate our choice of network architectures and also indicate with which baseline methods they are shared.
Deterministic encoder (ROTOC, Proto- RL, DB, RPC). We employ the convolutional encoder archi- tecture from Proto-RI. (Yarats et al. 2021a to parametrize both the online and target deterministic encoders, 9, and om. Each deterministic encoder consists of four convolutional layers following a single fully connected layer with Couv(filee==2,, kernel-size=3, stride=2) - Conv Cfnv(fil kernel- size=3, stride=1) - Conv(flter=32, kernel-size-3, stride=1) Conv(flter=32, kernel-size-3, stride=1) FCN(units=50) architocture. The ReLU activation function is used after cach con- volutional layer.
Stochastic encoder (ROTOC, RPC, DB). The architecture of the stochastic encoder is based on the state encoder from RPC [Eysenbach et al.) 2021 with modified hidden dimension. We change the hidden dimension from 256 to larger 1024 since ROTOC focuses on extracting embeddings from high-dimensional images rather than low-dimensional proprioception states in RPC. Specifi- cally, the online stochastic encoder Io is parameterized by a 2-layer fully-connected network with FCN(unis==0024) - FCN(units=100) architecture and ReLU hidden activations. Its output is divided into the mean 14, E 1 and the standard deviation &. R50 of the diagonal Gaussian distribution (=((cc))) The target stochastic encoder 9m shares the same network architecture as the online stochastic encoder.
Transition model (ROTOC, RPC, DB). We employ the architecture of the transition model from CoDy (You et al. 2022) with one minor difference to parametrize the transition model. Namely, to match the output dimension of the stochastic encoder, we change the output dimension from 50 to 100 for outputting the mean and standard deviation vectors of a diagonal Gaussian distribution. Specifically, the transition model qu consists of three fully connected layers with (CC(units=1024) + FCN(units=1024) + FCN(unis==000) architecture and ReLU hidden activations. Its output is divided into the mean ps R and standard deviation o E R500 of the diagonal Gaussian distribution (= a.).
Projection head (ROTOC, Proto-RL). The projection head d, is just a single linear layer which maps the concatenated bottleneck variable E R and action RIAI into a 50-dimensional vector. The projection head is based on the projector from Proto-RL, but accounts for the different imput dimensions.
Prediction (ROTOC,,,C, Proto-RL). The architecture of prediction heads is based on the predic- tor from Proto-RL. (Yarats et al.] (2211a) with modified hidden dimension. The online and target prediction heads, fo and fms both consist of two fully-connected layers with FCN(units=1024) FCN(units=50) architecture and ReLU hidden activations.
18

---

Published in Transactions on Machine Learning Research (08/2025)
B.3 SAC Architecture (ROTOC, RPC, ...)
We employ the publicly released standard Pytorch implementation (Yrats et al. 2021b) of SAC. We all hyperparameters of SAC from (Yarats et al. 202Tb), except for the replay buffer capacity, which we set to smaller 10 by following common practice in image-based RL (Yarats et al. 2021a You ct al. 2022). A smaller buffer can decrease memory costs for storing image observations. All hyperparaneterrs of SAC are fixed across tasks and shown in Table B.1 of the implementation of SAC.
We refer to Yarats et al. 2021b) for more detailed descriptions
Table B.1: Shared hyperparameters across tasks
Parameter
Value
Replay buffer capacity Optimizer
100 000
Adam
Critic Learning rate
10
Critic Q-function EMA
0.01
Critic target update freq Actor learning rate Actor update frequency
10-
Actor log stddev bounds
[-10 2
Temperature learning rate
10-3
Initial steps Discount
1000
0.99
Initial temperature
0.1
Learning rate for do. Jos 9a d, and fo Encoder and projection model EMA т
10
0.05
Coefficient Coefficient X Chunk length
0.1
0.001
B.4 The Deepmind Control Suite Setting
We evaluate our method ROTOC on set of challenging continuous control tasks from the commonly used benchmark Deepmind Control Suite. We set the episode length of each task to 1000 steps. Following PlaNet (Hafner et al. 2019b), we treat the number of action repeats as a hyperparameter of the agent. We use the number of action repeats from PlaNet (Hafner et al. 2019b for each task, which shown in Table
Table B.2: Task-spocific hyperparameters
Task
Action Repeats
Batchsize
Ball-in-cup Catch
256
Cartpole Swingup-sparse
256
Reacher Easy
256
Cartpole Swingup
8
256
Walker Walk
2
128
Cheetah Run
4
256
B.5 Other Hyperparameters
Following Proto-RL, the learning rate for optimizing networks Por 9or Go. d, and fo to 10-4, and the coefficient the exponential moving average used for updating target networks Φms 9m and fm set 0.05. Following CoDy (You et al. 2022). treat. batch size hyperparameter to the agent. We tested
19

---

Published in Transactions on Machine Learning Research (08/2025)
ROTOC with two batch sizes, 128 and 256, for each task. Based on empirical results, use batch size of 256 for all lasks, except for the Walker Walk task, which uses batch size of 128 (see Table A.2).
There are three specific hyperparameters of ROTOC which are chosen by performing hyperparameter tuning on standard Cartpole Swingup task and subsequently fixed for all tasks, namely coefficient a and À as well as the chunk length L. We provide a comprehensive overview of all hyperparameters of ROTOC in Table B.1
and Table B.2
B.6 InfoNCE Bound Implementation
In practice, we randomly sample a minibatch of (sequential) transitions (&t, 0t: 8t41), and obtain a minibatch of positive samples (zt, at, 2441) by encoding the minibatch using our modcls. For a given positive sample, the negative sample set *[41 is constructed by using all other embeddings 1141 of the same mini-batch.
B.7 Baseline Implementation
Proto-RL We obtain the results for Proto-RL by performing the original implementation provided by (Yarats et al.] 2021a) with one difference. We do not allow lask-agnostic pre-training of Proto-RL to facilitate fair comparison to setup. Instead, we perform the gradient updates of the repre- sentations and SAC agent from the first step of training by jointly optimizing Proto-RL objectives (LssL and LRL) given the task information.
DB We adapt DB to the off-policy continuous control setting. Similar to ROTOC, we employ the convolutional encoder architecture from Proto-RL (Yarats et al. 2021a) to parametrize the observation encoder for DB. The representation posterior and the prediction head of DB are based on the transition model from CoDy (You et al. 2022) with modified output dimension 100. To achieve as good performance and 03 in the information bottleneck objective LoB of DB and an additional hyperparameter B which augments DB bonus with environment rewards.
possible for DB, perform hyperparameter tuning to select suitable 0g
RPC As RPC do not provide an official implementation for image- based control lasks, we implement by ourselves by following (Eysenbach et al. 2021) as close as possible. Similar to ROTOC, we also employ the convolutional encoder architecture from Proto-RL (Yarats et al. 2021a) following the original stochastic state encoder architecture from RPC to parametrize pixel observation encoder for RPC. The dynamic model of RPC is based on the transition model from CoDy (You et al. 2022) with modified output dimension 100. We select the KI. constraint for RPC by hyperparameter tuning to achieve good performance. Specifically, we define a set of the KL constraint values, 31.0,3.0, 5.0]. and perform a grid search over it on the Cartpole Swingup task, and find the optimal KL constraint for RPC 3.0.
Plan2explore We the official implementation provided by (Pahhak el al. 2017 to acquire the results for Plan2explore.
To ensure a fair comparison, we use the above mentioned implementation of SAC for all methods in our experiments, except for Plan2Fxplore which follows original implementation and learns a policy basod on model-based RL method Dreamer (Hafner al., 2019a). The performance of each algorithm evaluated by computing average return over 10 episodes every 10K environment steps. For each method, we perform one gradient update per environment step for all tasks to ensure fair comparison. Notably, fixed chunk length L = 2 for our method across all tasks and the number of samples we used per gradient update step is al always twice the batch size. We set the batch size to 512 for Proto-RL by following the official implementation (Yarats et al. 2021a). We set the batch size to 512 for DB and RPC by following Proto-RU RL and CURL. For Plan2explore, we use the amount of data per gradient step from its official implementations. Our method uses the same amount of samples per gradient step as all model-free baselines on the majority of tasks, except for the Walker Walk task where
use the
smaller number of samples.
20

---

Published in Transactions on Machine Learning Research (08/2025)
B.8
Robustness to Noisy Observations
For this experiment, we add pixel-level Gaussian white noise to the background of each rendered image. The main body of the manipulated object in each task remains undisturbed for efficient observation.
B.9
Robustness to Natural Background
Following [Zhang et al. 2020 the background of the Mujoco tasks replaced with by videos randomly sampled from the driving car class in the Kinetics datasct. The backgrounds of the Mujoco tasks are constantly changing during training and evaluation.
B.10 Algorithm
The training procedure of MTC is presented in Algorithm 1 new experiences from the environment, and updating the parameters of the ROTOC model, the actor and critic networks of SAC.
The algorithm alternates between collecting
Training Algorithm for ROTOC Initialize: The ROTOC model, actor and critic networks, replay buffer D, Batch size B, Chunk length L
Algorithm 1:
for each training step do collect experience (st, for each gradient step do
and add the replay buffer
Sample minibatch of sequence chunks {( ai rí, (111) 3iml D from replay buffer. Generate {(c), c+1, 21) L+k-11 - following Eq
L+k
Compute augmented reward laugg (8,, ar) by following Eq,
Update the online deterministic and stochastic encoders, the latent transition model, the projection head, the online prediction model and the weight transformation matrix by following
Eq. 7
Update the actor and critic networks of SAC.
Update the target deterministic and stochastic encoders and the target prediction head.
end
end
21

---

Published in Transactions on Machine Learning Research (08/2025)
C
Additional Performance and Robustness Evaluation
Fig.
and Fig. C.2 show the performance of our method and baselines on noisy and natural Mujoco tasks,
respectively. Fig. C.3
shows the performance our method and baselines on all six Mujoco tasks in all three settings to see the performance drops caused by the distractive information.
Cartpole Swingup
Ball-In-cup Catch
Reacher Easy
1000
1000
Eolsode reward
Episode reward
Episode reward
800
600
400
200
0.0
0.1
6.2
0.4
0.5
0.0
0.1
6.2
0.4
0.0
6.1
0.2
Cheetah
Cartpole Swingup_sen _parress__reeee_
400
1000
800
Eaisode reward
Episode reward
800
Eolsode reward
300
600
200
400
400
200
200
0.4
0.6
0.1
0.2
0.3
0.2
Envinronment stepsix109)
Envinronment stepsix109)
Envinronment stepsix109)
ROTOC
Plan2Explore
RPC
Figure C.1: Robustness comparisons six noisy Mujoco tasks. compared to the baselines.
Our method performs bost on all tasks
Cartpole Swingup
Ball-in-cup Catch
Reacher Easy
1000
1000
Gsode reward
Episode reward
Episode reward
500
600
600
200
200
200
Cheetah Run
Walker Walk
Cartpole Swingup_ spparss
Eolsode reward
Episode reward
Eolsode reward
200
100
200
0.4
0.4
Envinronment stepsix10
Envinronment stepsix10*)
Envinronment stepsix109
ROTOC
DB
Proto-RL
Plan2Explore
RPC
Figure C.2: Performance comparisons on six natural Mujoco lasks. Our method outperforms other hods in robustness to natural video backgrounds.
We also compare the performance achieved by our method and baselines 500K environment steps on the noisy natural tasks with the final performance evaluated the standard tasks for directly checking the
22

---

Published in Transactions on Machine Learning Research (08/2025)
Cartpole Swingup
Ball-in-cup Catch
Reacher Easy
1000
Eplsode rewaeeaaa sward
Episode reward
Episode reward
600
400
200
200
0.0
0.1
6.2
6.3
0.4
0.0
0.1
6.2
0.3
0.4
6.5
0.0
0.2
0.4
0.6
0.0
Cheetah Run
Walker Walk
Cartpole Swingup_sse s_paasee
800
co.sode reward
200
Episode reward
Eelsode reward
600
800
400
400
200
0-
0-
0.4
0.6
0.2
0.3
0.4
0.6
Envinronment steps(x109)
Envinronment stepsix109)
Envinronment steps(x109)
ROTOC
DB
Proto-RL
Plan2explore
RPC
ROTOC (Noise)
DB (Noise)
Proto-RL (Noise)
2Ixn2Ex (Noise)
RPC (Noise)
ROTOC (Video) DB (Video)
Proto-RL (Video)
Plan2Explore (Video)
RPC (Video)
We show the performance of method and baselines all six Mujoco tasks in all three settings to directly compare the performance drops caused by the presence of distractive information ( (both white noise and natural vidco backgrounds).
Figure C.3:
Table C.1: Scores achieved by our method (mean and standard error for 10 sceds) and baselines at 500k environment steps on noisy Mujoco tasks and natural Mujoco tasks, with the bracketed scores achieved by the corresponding methods 500K environment steps or the last values we evaluated on standard Mujoco tasks. The dagger (f) indicates that the scores in parentheses
achieved at less than 500K environmental steps. The bold font indicates that the lower-bound the reward (based the given standard error intervals) for the given method, is larger or equal Ihan the respective upper bound for every other exploration method.
500K step Scotes
OOTS
DB
Proto-RL (823
Plan2Explore
RPC
Cartpole Swingup
(870 3
(800
773228 [51l 1
814213
Ball-in-cup Catch
964 + (957
3"
182 + 155 (067 + 5")
9224 16 (913 261)
802-443 (904 + 501)
928-200 (947 8")
Noisy
Mujoco
Reacher Easy
917: 20 (969
14)
939 + (927 36)
869: (896
76)
8511 132 (892 + 6)
930: 29 (949
20)
Chestah Run
332+
(570
79)
265
(577
208+ 13
20)
2514 (430 25)
282+18 (498
74)
tnoks
Walker Walk
935 11 (952
+ 3
886 + 30 (687 = 120")
9154t 15 (885 = -11)
752=36 (646 = 72")
893=11 (858
:: 411)
Cartpolo Swingup-spuren
(793
10)
133 (653 94)
571+ 95 (757+ 23)
680+ 51 (067
724
(177
107)
Cartpole Swingup
8441 (870
31)
503 16 (860 1.5")
7481 28 (62) 24"
143122 (541 1 63
812 (861
Ball-in-cup Catch
917:2 12
(957
+ 3
916 + 17 (957 + 5f)
7274 53 (913 261)
133:136 (901 50f)
838-722 (941 = 8")
Natural
Mujeco
Reacher Easy
14)
(927
(896 76)
502:1
tasks
Cheetah Run
308: 16
(570
= 79)
294 + 16 (G77 91)
2964 (631 20)
1934 32 (130 25)
279:12 (987 + 120")
Walkerr Walk
886:1 15
[952
1 3
841 1 (687 120
882:1 20 (885 = 41
557254 (546 : 72
714:26 (858 41
Cartpole Spiagup-sparse
10)
51)
5 (177 107)
performance drops. Notably, we just evaluated sample efficiency in the small data regime on some standard tasks, which was enough for our method to converge, but not enough for some of the baselines. In Table C.I we indicate when the scores are computed at different numbers of stops to avoid wrong conclusions.
23

---

Published in Transactions on Machine Learning Research (08/2025)
D Ablation Studies
D.1 Investigating Chunk Lengths
To evaluate the future predictability of the learned stochastic representations, we first learn representations on the Cartpole Swingup-sparse task with natural backgrounds, and then fix the representations. While using different chunk lengths, we use the same amount of data for sampling from the replay buffer per gradient step. For the fixod stochastic emboddings, we build a forward prediction model parameterized by neural network, which takes the fixed embedding of the current stale at time step 1, ãt; and outputs the mean and the standard deviation of the diagonal Gaussian distribution d(zztn/z) with prediction step We set the prediction step to 1, 2, 3, or 5. The prediction model consists of three fully connected layers with 5CN(mits=512) + 5C) + train the prediction model by maximun-likelhhod estimation. We did the same for evaluating the future predictability of the learned deterministic representations.
architecture and ReLU hidden activations. We
Table D.1 compares the future predictabllity of the learned deterministic representations with different chunk lengths on the Cartpole Swingup-sparse task. Compared to using individual transitions, the learned reprosentations by using larger chunk lengths L 2,3 achieve better prediction at distant prediction steps 3,5, while achieving worse prediction at prediction steps 1, This indicates that maximizing objective doesn't improve the one-step correlation between latent deterministic embeddings, but causes an increase in the long-term correlations.
Table D.1: We test the future predictability of the learned deterministic reprosentations by optimizing a multi-step predictive model on a data set of latent representations. We obtain larger log-likelihoods with prediction step = 3, 5 and smaller log-likelihoodd with prediction step = 1,2 when using an encoder that was trained with chunk length L 2,3, compared to using an encoder with chunk length L
Prediction Steps
3
1I=
19.2 + 0.3
14.6
0.1
7.6 + 0.1
0.2
L=1,BS=512
22.8 + 0.2
15.1 - 0.1
6.4 0.1
-2.4 + 0.1
L=2,BS=256
17.0 + 0.3
14. 0.2
6.9 0.0
3.7 0.2
L=3,BS=256
16.0 1 0.2
14.31 0.1
8.6 2 0.1
4.6 0.2
L=5,B8=256
0.5 = 0.3
-4.2 0.1
-11.6
-18.4 + 0.4
D.2 Full Results on Three Reward-sparse Tasks
We show the performance ROTOC and its three ablations, and vanilla SAC on three reward-sparse tasks in three settings in Fig. D.I] ROTOC achieves better least comparable performance its own ablations across all tasks. Moreover, ROTOC significantly outperforms vanilla SAC across all tasks.
D.3 Effect of KL-based Bound on Improving Predictability
We also compare the future predictability of the learned stochastic and deterministic representations with without the KL-based lower bound the standard Cartpole Swingup-sparse task. The results in Table D.2 show that the learned stochastic and deterministic representations by ROTOC are better suited for future prediction compared to the representations learned by Non-KL. ROTOC that lacks the KT.-based lower bound, indicating that using the KL-based lower bound improves the correlations between the learned representations.
D.4 Investigating Intrinsic Rewards
D.4.1 Intrinsic Rewards on Two Subsets of States
Fig. D.2 shows the initial cart position and pole angle in the Cartpole Swingup-sparse task with random natural background. A pole angle in the range of [355", 55]] corresponds to upright pole position, while
24

---

Published in Transactions on Machine Learning Research (08/2025)
Cartpole Swingup-sparse
Bell coll-in-cup Catch
Reacher Easy
1030
030
Episode reward
800-
600
Egiisde
600
Episede
200
420
200
5.0
0.0
0.05
0.15
5:35
0.22
0.
62
Envinronment stepsix10*)
Envinronment steps(x10*)
Enviinonmmnt steps(x100)
Cartpole Swingup sparse (White Noise)
Ball in- Catch (White Noiso)
Reacher Easy (White Noise)
1000
Epicade reward
820
Episode reward
800
600
600-
800
Episode
200
270
6.
03
Envinronment stepsix109)
Envinronment steps(x10*)
Envinronment steps(x10*)
Cartpole Swingup-sparse (Natural Video)
(Natural Video)
Reacher Easy (Natural Videa)
Ggioode reward
030
Episode reward
E00
600
Episode
200
200-
Envinronment stepsix109)
Envinrenment steps(x110)
Envinronment steps(x110))
62
ROTOC
No No-InfoNCE ROTOC
No KL- ROTOC
No Reward ROTOC
Vanillaa SAC
Figure D.1: Performance of ablations on three standard Mujoco tasks (top row), noisy Mujoco tasks (middle row) and natural Mujoco tasks (bottom row). Each subplot shows the average reward for seeds with 90% confidence intervals shading.
Table D.2: The future predictability the representations learned by ROTOC and Non-KT, ROTOC, respectively. The stochastic and deterministic representations learned with the KL bound are better suited for future prediction.
Prediction Steps
5
ROTOC
-80.9
-91.9 1 0.4
-99.6
-0.1
-109.3 1: 0.4
Stochastic Representation
Non-KL ROTOC
-86.5 + 0.3
-94.3t 0.3
-10066 + 0.2
-109.8+ 0.2
ROTOC
17.0 + 0.3
14.4 = 0.2
6.9 0.0
3.7 0.2
Deterministic Representation
Non-KL ROTOC
9.6 0.1
7.4+ 0.2
3.3 0.3
0.3 + 0.1
angle the range of [175", 185"] corresponds a pole pointing downwards. The cart position in the range of [-0.25, 0.25] corresponds to middle position.
To investigate the effectiveness of our intrinsic reward function, save the intrinsie reward function every 100,000 steps during training. We then save the replay buffer after training our model from scratch on the task to ensure that the oblained replay buffer is independent of the saved reward function. We sampled a fixed data set from the replay buffer, and split the data set into two, where one of the subsets only contains states with a pole angle in the range of 355",5" and the other subset only contains states with a pole angle in the range of [175', 185"]. For each subset, we sample a minibatch of states and compute the intrinsic rewards of the drawn samples by using the saved reward function.
25

---

Published in Transactions on Machine Learning Research (08/2025)
0.25
185
Figure D.2: The pole is initialized pointing downwards with the pole angle 180 and located in the middle with the position 0 in the Cartpole Swingup-sparse task.
D.4.2 Intrinsic Rewards on Trajectories
We also compute the intrinsic rewards along trajectories sampled from the saved replay buffer. We sample trajectory and then compute the intrinsic rewards of this trajectory by using the reward functions saved the different phases of the training. Specifically, we use the reward function saved at 200K enviromment steps where the policy achieves a reward of around 0 (see Fig. C.2), the reward function saved at 600K environment steps where the reward achieved by the policy is increasing rapidly, and the reward function saved at 1000K environment steps where the policy becomes good at balancing the pole. Fig. D.3 Fig. D.4 and Fig. [.5 show the evolution of the intrinsic reward along an ncar-optimal trajectory. Regardless of the distractive backgrounds, the reward functions at 200K and 600K environment steps tend to allocate higher intrinsic rewards to images that show upright pole position, than images that show pole pointing downwards. The reward function saved at 1000K steps does not show this tendency anymore, since these images that show an upright pendulum become more common in the replay buffer.
Trajectory Intrinsic reward
Reward
Reward function
Intrisic reward
0.6
0.4
0.2
Timestep
Figure D.3: Evolution of intrinsic rewards along the first of three near-optimal trajectorics sampled from the replay buffer for the CartPole Swingup-sparse task. The dots along the x-axis correspond to observations in the sampled trajectory per 80 time steps. Initially (blue curve), higher rewards are assigned to states with an upright pendulum, which are rarely observed. As training progresses (orange and green curves), the intrinsic rewards for such states decrease as they become less novel.
26

---

Published in Transactions on Machine Learning Research (08/2025)
Trajectory Intrinsic reward
1.0
Reward steps
Reward function 6O0K steps
Ranward function 1000K staps
0.8
Intiisic revard
0.6
0.0
Timestep
Figure D.4: Evolution of intrinsic rewards along the second of three near-optimal trajectories sampled from the replay buffer for the CartPole Swingup-sparse task. The dots along the x-axis correspond to observations in the sampled trajectory per 80 time steps. Initially (blue curve), higher rewards are assigned to states with upright pendulum, which are rarely observed. As training progresses (orange and green curves), the intrinsic rewards for such states decrease they become less novel.
Trajectory Intrinsic reward
1.0
Reward function at 2008 stops
Revard function at BOOK steps
Reward function at 1000K steps
0.8
Intiraic reward
0.6
0.0
Timestep
Figure D.5: Evolution of intrinsic rewards along the second of three near-optimal trajectorics sampled from the replay buffer for the CartPole Swingup-sparse task. The dots along the x-axis correspond to observations in the sampled trajectory per 80 time steps. Initially (blue curve), higher rewards are assigned to states with upright pendulum, which rarely observed. As training progresses (orange and green curves), the intrinsic rewards for such states decrease they become less novel.
27

---

Published in Transactions on Machine Learning Research (08/2025)
D.5 Investigating Effect of Chunk-wise sampling Optimization
In this section, we perform an ablation experiment to invesligate whether using chunk-wise mini-batching sampling produces more consistent gradients along trajectories. In our experiment, on the Cartpole Swingup- sparse task, we collected 20 independent fixed trajectories, computed the gradients the parameters of the encoder with respect to the loss every time step in trajectories, and then calculated the variance of gradients over all time steps during training. We compare the curves of the gradient variance for different. chunk lengths in Fig. [.6 and the final gradient variance in Table D.3 While using the same number of data samples, using chunking with L 2 2 achieves smaller gradient variance than using chunk length of 1. We also observed that the variance of using individual transitions significantly increases faster than using larger transitions when the number of samples increases. These results demonstrate that using chunk-wise mini-batching sampling makes the gradients along a trajectory more consistent, helping learning globally consistent representations.
512 Samples
768 Samples
1280 Samples
100
B5=512
B6= L=3, BS=256
BS=1280
Gradient Variance
120
L=2, B5=256
5, BS=256
100
60
40-
0.35
0.4
0.45
0.5
20
0-1
0.0
0.1
0.2
0.3
0.4
0.5
0.0
0.1
0.2
0.3
0.4
0.5
0.0
0.1
0.2
0.3
0.4
0.5
Envinronment Step(x10")
Envinronment (tep(x10")
Figure D.6: We compare the variance of the gradients over time steps for different chunk length L and batch size BS. The plot shows the mean and standard error for 3 seeds. Using chunking produces a smaller gradient variance along a trajectory than without chunking, when using the same number of samples.
Table D.3: We compare the variance of the gradients over time steps for different chunk length L and batch size BS. The table shows the mean and standard error for 3 seeds. When using the same amount of training samples, using sequential transitions achieves smaller gradient variance than using individual transitions (L=1).
512 samples
768 samples
1280 samples
Configurations
L-1, 512
L-2,
BS-256
L-1,
BS-768
L-3, BS-256
L=1, BS USS=122 1280
L L=5, BS-256
Gradient variance
3.54 2.0
3.1 1.8
10.9 6.3
5.91 3.3
28.74 16.6
3.2 1.8
28

---

Published in Transactions on Machine Learning Research (08/2025)
E
Model Complexity
We evaluate the model complexity of ROTOC and other methods. Table E.4 shows the number of learnable parameters for all methods. The model complexity of ROTOC is higher than RPC, but is lower than DB and Proto-RI.
Table E.4: Comparison of model complexity. ROTOC has 3.5 million learnable parameters, slightly higher than DB (3.4 million) and lower than DB and Proto-RL.
Method
ROTOC (Ours)
RPC
DB
Proto-RL
Number of Parameters
3.5M
3.4M
4.6M
5.2M
29

---

Published in Transactions on Machine Learning Research (08/2025)
F Representation Visualization
Fig. Fig. and Fig. visualize the learned representations by RPC, Proto-RI. and DB on the natural Walker Walk task, respectively. We also visualize the learned deterministic representation learned by our method on the natural Walker Walk task in Fig F.4 We evaluate learned representations by seeing if representations that have similar robot configurations and clearly different backgrounds appear close to each other. The representations learned by Proto-RL and RPC look worse, but the representations learned by DB and deterministic embeddings comparable to the stochastic embedding learned by method.
Figure F.1: The representations learned by RPC visualized using t-SNE 20 x 15 grid.
30

---

Published in Transactions on Machine Learning Research (08/2025)
Figure F.2: The representations learned by Proto-RL are visualized using t-SNE on 20 15 grid.
Figure F.3; The representations learned by DB are visualized using t-SNE on a 20 15 grid.
31

---

Published in Transactions on Machine Learning Research (08/2025)
Figure F.4: The deterministic representation learned by our method on the Walker Walk task with natural background are visualized using t-SNE on a 20 x 15 grid.
32


Reviews:
[{"review_id": "CTHYeTbTuf", "reviewer": "Reviewer_7GyU", "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\n1.  **Novel and Intuitive Objective**:\n    * The core **Rollout Total Correlation (ROTOC)** objective, $\\max \\mathcal{C}(z_1; a_1; \\dots; a_{T-1}; z_T)$, is a well-motivated and principled approach to learn representations that capture long-term temporal consistency. This directly addresses limitations of prior methods focusing on individual transitions.\n2.  **Comprehensive and Tractable Method**:\n    * The paper successfully translates the intractable ROTOC objective into a practical algorithm by deriving a lower bound based on the sum of per-step mutual information, $\\sum_{t=1}^{T-1} I(z_{t+1}; z_t, a_t)$.\n    * It employs two complementary lower bounds for the per-step mutual information: a generative (KL-divergence based) model for consistency and a discriminative (InfoNCE based) model for representativeness.\n    * The introduction of **chunk-wise mini-batching** is a simple yet effective technique to encourage the learning of longer-term correlations, helping to align the per-step optimization with the global objective.\n    * An **intrinsic reward** derived from the InfoNCE loss is integrated to encourage task-specific novelty and exploration.\n3.  **Strong Empirical Results and Robustness**:\n    * The method demonstrates improved sample efficiency and often superior asymptotic performance compared to several leading baselines on a set of challenging image-based Mujoco control tasks.\n    * ROTOC shows significant robustness to observational distractors, including Gaussian white noise and dynamically changing natural video backgrounds, often outperforming baselines in these settings.\n4.  **Thorough Ablation Studies and Analysis**:\n    * Extensive ablation studies systematically validate the contributions of different components: chunk-wise mini-batching, the individual lower bounds (KL and InfoNCE), and the intrinsic reward.\n    * The paper provides evidence that chunk-wise mini-batching increases multi-step predictability of the learned representations, linking this to improved performance.\n    * Visualizations (e.g., t-SNE plots) suggest that ROTOC learns representations that effectively capture task-relevant features (like robot configurations) while being invariant to background distractors.\n\n\n## Weaknesses\n\n1.  **Approximation of the Core Objective**:\n    * While the ROTOC objective $\\max \\mathcal{C}(z_1; \\dots; z_T)$ targets long-term correlation, the primary tractable lower bound decomposes into a sum of per-step mutual information terms $\\sum I(z_{t+1}; z_t, a_t)$. As the authors acknowledge, this decomposition might inherently lose some of the direct incentive to capture long-term correlations beyond single steps.\n    * Chunk-wise mini-batching is a practical and empirically effective technique to mitigate this, but it's an indirect way to optimize the original global objective, as the loss function itself remains a sum of per-step terms.\n2.  **Potential for More Direct Multi-Step Modeling**:\n    * The authors themselves suggest that \"it would also be interesting to explore objectives that better capture multi-step correlations based on multi-step prediction models\". This indicates an area where the current approach could be extended for even stronger long-term dependency modeling, moving beyond the current one-step latent transition model $q_{\\psi}(z_{t+1}|z_t, a_t)$.\n3.  **Clarity on Negative Sample Selection for InfoNCE**:\n    * For the InfoNCE lower bound, negative samples $z_{t+1}^*$ are drawn from the marginal distribution $p(z_{t+1})$. While standard, practical implementation often involves sampling negatives from the current batch or a memory bank. Further details on how these negative samples are specifically obtained/approximated in the experiments could enhance reproducibility", "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": "The authors study the problem of learning task-relevant representations in image-based reinforcement learning, with a focus on long-term temporal consistency and objectives that incentivize consistency over long horizons.\n\n1.  **ROTOC Objective**:\n    * Proposes maximizing the total correlation among latent state embeddings ($z_t$) and actions ($a_t$) within entire trajectories: $\\max \\mathcal{C}(z_1; a_1; \\dots; a_{T-1}; z_T)$.\n\n2.  **Tractable Optimization Framework**:\n    * The ROTOC objective is lower-bounded by the sum of per-step mutual information: $\\sum_{t=1}^{T-1} I(z_{t+1}; z_t, a_t)$. \n    * This per-step $I(z_{t+1}; z_t, a_t)$ is further approximated by combining two complementary lower bounds:\n        * A **generative model-based bound** using KL divergence (for predictable and consistent representations) .\n        * A **discriminative model-based bound** using InfoNCE loss (avoid collapse and trivial solutions of KL bound): $I_{\\omega}(z_{t+1};z_{t},a_{t})$.\n    * The final representation learning loss combines these terms:\n        $L = \\mathbb{E}[\\sum_{t=1}^{T-1}\\alpha$ $\\mathbb{D}\\_{KL} ( g_m $(z\\_{t+1} $ $|c\\_{t+1}) || q\\_{\\psi} (z\\_{t+1}|z_t,a_t))$ $ - $I_{\\omega}(z_{t+1};z_t,a_t)]$.\n\n3.  **Chunk-Wise Mini-Batching**:\n    * A technique where mini-batches are sampled as sequence chunks $(s_t, a_t, s_{t+1})_{t=k}^{L+k-1}$ of length $L$.\n    * This is proposed to better approximate the original ROTOC objective and encourage learning of long-term correlations, despite the per-step decomposition of the bound.\n    * Empirically shown to increase long-term predictability and agent performance.\n\n4.  **Intrinsic Reward for Exploration**:\n    * An intrinsic reward $r^*(s_t, a_t)$ is derived from the InfoNCE loss term to promote exploration.\n    * Specifically, $r^*(s\\_t , a\\_t ) = $ $ -\\mathbb{E}_{p(z\\_t, a\\_t, z\\_{t+1})} $ $[\\text{InfoNCE log term} ]$ , with higher values for less predictable transitions.\n\n5.  **New Knowledge from Empirical Validation**:\n    * ROTOC achieves improved sample efficiency and final performance on standard image-based Mujoco tasks.\n    * Demonstrates enhanced robustness to observational noise (Gaussian white noise and natural video backgrounds) compared to baselines.\n    * Ablation studies confirm:\n        * Chunk-wise mini-batching (e.g., $L=2,3$) outperforms individual transitions ($L=1$) and leads to better multi-step predictability of representations.\n        * Both the KL-based and InfoNCE-based lower bounds, as well as the intrinsic reward, contribute positively to performance. The KL-based bound specifically improves long-term predictability.\n        * Learned representations effectively filter out distractors and capture task-relevant features (e.g., robot configurations).", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "1.  **Clarification of the Link Between Per-Step Loss and Global ROTOC Objective with Chunk-Wise Mini-Batching**\n    * **Specification**: *Would significantly clarify and strengthen the work's claims regarding the optimization of long-term correlations.*\n    * **Adjustment**: Elaborate further in Sections 4.3 and 4.5.1 on how optimizing the sum of per-step mutual information objectives, $\\sum_{t=1}^{T-1}I(z_{t+1};z_{t},a_{t})$, using chunk-wise mini-batching helps in achieving the global objective of maximizing rollout total correlation $\\mathcal{C}(z_{1};a_{1};\\cdot\\cdot\\cdot;a_{T-1};z_{T})$. The current text states chunk-wise mini-batching \"may help to converge to local optima with better total correlation\" and relies on the intuition that \"a stochastic gradient that increases correlation between several subsequent representations [will] favor temporally consistent representations\". Expanding on *why* this batching strategy, in conjunction with the per-step loss, is hypothesized to better approximate the optimization of the original long-term objective would be beneficial. For instance, discuss if it encourages smoother representation manifolds over sequences or if the aggregated gradients from chunks better reflect the dependencies across the entire rollout.\n\n2.  **Explicit Details on InfoNCE Negative Sample Strategy**\n    * **Specification**: *Would strengthen the work by enhancing reproducibility.*\n    * **Adjustment**: In Section 4.4.2 or Appendix B, provide precise details on how the negative samples $z_{t+1}^{*}$ (for the InfoNCE loss $I_{\\omega}(z_{t+1};z_{t},a_{t})$) are obtained in practice. The text mentions they are \"drawn from the marginal distribution $p(z_{t+1})$\", but it would be helpful to clarify if these are, for example, other $z_{t+1}$ samples from the same mini-batch, from a separate memory bank, or another common approximation strategy.\n\n3.  **Brief Acknowledgment of Hyperparameter Tuning Scope for Representation Learning Components**\n    * **Specification**: *Would strengthen the work by improving transparency regarding the experimental methodology.*\n    * **Adjustment**: Add a sentence in the experimental setup (Section 5.1) or implementation details to explicitly state that the primary representation learning hyperparameters (e.g., coefficient $\\alpha$, $\\lambda$ for the intrinsic reward, and chunk length $L$) were selected based on performance on a specific task (standard Cartpole Swingup) and then applied consistently across all other tasks and environments.\n    * **Reasoning**: This highlights the generalization of the chosen hyperparameter set while implicitly noting that task-specific tuning wasn't exhaustively performed for every environment, providing a complete picture of the tuning strategy.", "broader_impact_concerns": "The research focuses on improving representation learning techniques for reinforcement learning agents, primarily validated in simulated control environments. The specific methodology of maximizing rollout total correlation does not appear to introduce unique ethical concerns beyond those generally applicable to AI and reinforcement learning research. The paper's aim to enhance robustness to visual distractors could be viewed as a positive step towards more reliable AI systems.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "J1jdAgHPyT", "reviewer": "Reviewer_dHTV", "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 contributions and motivation are clearly stated and (relatively) well justified.\n- The paper is clear, and the experiments are sufficient to back the conclusions and contribution statements. \n- The main idea is intuitively straight-forward.\n\n## Weaknesses:\n- I'm not entirely convinced by the motivation (although it is clear, as mentioned). While considering long-term information gain in sequential problems for learning better representations seems reasonable, the reasons for why this should induce an across-the-board improvement of RL algorithms for robotic tasks is not clear to me, and in fact does not seem to be the case from the experimental results (there are clear improvements in some tasks, but it is not obvious to me that these are not due to task selection bias).\n- The resulting lower bounds (which make the objective tractable) seem to result in a loss which is essentially reduced to a sum of mutual information between consecutive time-steps. For this reason, the main difference with respect to e.g. RPC (which proposes a mutual information penalty to learn more compact and predictable representations in off-policy RL) seems to be the mini-batch training heuristic, where the data for training is sampled in batches of consecutive trajectory points to bias the learning towards 'longer term' consistency. This seems like a slightly thin contribution given the body of existing work.", "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": "The authors present a method for increasing sample efficiency in reinforcement learning through incorporating temporally consistent representation learning in an off-policy RL setting. They argue that optimizing for total correlation of the state action sequences is beneficial for learning efficient representation of dynamic tasks from image data, propose an encoder architecture and an information theoretic lower bound for such total correlation based on mutual information estimates, incorporate this into a proposed modified reward function and apply SAC with such modified reward function. They demonstrate in image-based robotic tasks how their method improves (in most cases) over baselines.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "1. Can the authors give a more detailed justification for why considering correlation objectives over entire trajectories would help RL policies perform better in robotic tasks (or in any case, learn faster/with less data)? What types of tasks, and for which tasks such objectives are not expected to induce any improvement?\n2. Can the authors provide more detail on the specific differences and novelty in their method with respect to other algorithms that already consider mutual information objectives in off-policy RL (e.g. RPC)?\n3. Given that RPC also uses SAC as a base implementation, and that it also computes MI estimates between latent state-action pairs and uses them to learn more compact representations, I would expect the empirical performance of RPC and ROTOC to be closer in most tasks. What is the reason for RPC having such poor performance in the experiments chosen when compared to ROTOC? What are the main differences in both implementations that cause these?", "broader_impact_concerns": "n/a", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "tcLbvO2ssC", "reviewer": "Reviewer_DXqP", "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:\n1. The objective is novel\n2. The empirical comparison is comprehensive and solid\n\n>Weakness:\n1. How close or tight is the proposed lower bound compared to the true rollout total correlation?\n2. since earlier states influence future ones. However, the best results are observed at L = 3. What factors limit the performance at higher chunk lengths? Additionally, will different environments have different best perform L?\n3. This method involves many NN models. Although the result is promising, how can we fine tune all the models? e.g. determine efficent network size for the models?\n4. How is the computaion complexity of the method with such complex models?", "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 proposed a Rollout Total Correlation (ROTOC) as a novel objective for representation learning in reinforcement learning, designed to capture long-term temporal dependencies within trajectories. This paper derive a lower bound by combining two complementary approximations of per-step mutual information: one generative and one discriminative. Additionally, they introduce a chunk-based mini-batching strategy to better preserve multi-step correlations during training.  The Results show that the method significantly improves sample efficiency and robustness to observation noise compared to leading baselines.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "Please refer to above weakness.", "broader_impact_concerns": "N/A", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Decision: Accept (poster)