Ego-Foresight: Self-supervised Learning of Agent-Aware Representations for Improved RL
Authors: Manuel Serra Nunes
Atabak Dehban
Yiannis Demiris
Jose Santos-Victor
Venue: ICLR 2026 Poster (2026)
Keywords: Agent-awareness
Reinforcement Learning
Self-supervised Learning
Robotics
Primary area: reinforcement learning
TL;DR: We present a self-supervised method for disentangling agent information in visual RL and show that the resulting features improve sample-efficiency and performance in existing model-free and model-based RL algorithms.

Abstract: Despite the significant advances in Deep Reinforcement Learning (RL) observed in the last decade, the amount of training experience necessary to learn effective policies remains one of the primary concerns in both simulated and real environments. Looking to solve this issue, previous work has shown that improved efficiency can be achieved by separately modeling the agent and environment, but usually requires a supervisory signal. In contrast to RL, humans can perfect a new skill from a small number of trials and often do so without a supervisory signal, making neuroscientific studies of human development a valuable source of inspiration for RL. In particular, we explore the idea of motor prediction, which states that humans develop an internal model of themselves and of the consequences that their motor commands have on the immediate sensory inputs. Our insight is that the movement of the agent provides a cue that allows the duality between the agent and environment to be learned.  To instantiate this idea, we present Ego-Foresight (EF), a self-supervised method for disentangling agent information based on motion and prediction. Our main finding is that, when used as an auxiliary task in feature learning, self-supervised agent-awareness improves the sample-efficiency and performance of the underlying RL algorithm. To test our approach, we study the ability of EF to predict agent movement and disentangle agent information. Then, we integrate EF with both model-free and model-based RL algorithms to solve simulated control tasks, showing improved sample-efficiency and performance.

Published as a conference paper at ICLR 2026
EGO-FORESIGHT: SELF-SUPERVISED LEARNING OF AGENT-AWARE REPRESENTATIONS FOR IMPROVED RL
Manuel Serra Nunes* Atabak Dehban', Yiannis Demiris", José Santos- Victor Institute for Systems and Robotics, Instituto Superior Técnico, U. Lisboa 2 Personal Robotics Laboratory, Imperial College London
ABSTRACT
Despite the significant advances in Deep Reinforcement Learning (RL) observed in the last decade, the amount of training experience necessary to learn effective policies remains one of the primary concerns in both simulated and real envi- ronments. Looking to solve this issue, previous work has shown that improved efficiency can be achieved by separately modeling the agent and environment, but usually requires a supervisory signal. In contrast to RL, humans can perfect a new skill from small number of trials and often do so without supervisory signal, making neuroscientific studies of human development valuable source of inspi- ration for RL. In particular, we explore the idea of motor prediction, which states that humans develop internal model of themselves and of the consequences that their motor commands have on the immediate sensory inputs. Our insight is that the movement of the agent provides and environment to be learned. To instantiate this idea, we present Ego-Foresight (EF), supersuu method for disentangling agent information based mo- tion and prediction. Our main finding is that, when used as an auxiliary task in feature learning, self-supervised agent-awareness improves the sample-effciency and performance of the underlying RL algorithm. To test our approach, we study the ability of EF predict agent movement and disentangle agent information. Then, we integrate EF with model-free and model-based RL algorithms to solve simulated control tasks, showing improved sample-efficiency and performance.
that allows the duality between the agent
INTRODUCTION
While it usually goes unnoticed go about our daily lives, the human brain is constantly engaged in predicting imminent future sensori inputs (Clark, 2016). This happens when we react to a friend extending their arm for handshake, when we notice missing note in our favorite song, and when we perceive movement in optical illusions (Watanabe et al., 2018). At more fundamental level, this process of predicting our sensations is seen by some as the driving force behind perception, action, and learning (Friston, 2010). But while predicting external phenomena is a daunting task for brain, motor prediction (Wolpert& Flanagan, 2001) i.e., predicting the sensori consequences of one's own movement- is remarkably simpler, yet crucial in increasing saliency of external events (Blakemore et al., 1999). Anyone who has ever tried to self-tickle has experienced the brain in its predictive endeavors: in trying to predict external (and more critical) inputs, the brain is thought to suppress self-generated sensations, to increase the saliency of those coming from outside- making hard to feel self-tickling (Blakemore al., 2000).
In artificial systems, a significant effort has been devoted to learning World Models (Ha & Schmid- huber, 2018; Finn et al., 2016; Gumbsch et al., 2024), which designed to predict future states of the whole environment and allow planning in the latent space. Despite encouraging deployments of these models in real-world robotic learning (Wu et al., 2023), their application remains restricted to safe and simplified workspaces, with sample-efticient Deep Reinforcement Learning (RL) being one of the main challenges.
"onresponning author: Email manueferammuns@rrciioo..iii.i.

---

Published as a conference paper at ICLR 2026
sealed by grad. scene feutures lt
Prod. scaled by grad. agent fentures Rgg
G.T.
At 30k toiin steps At 10k: train steps
BRO 2F
DRQQ2 22
SEAR
Matverdito
BRO12-2
15 SS
SEAA
Figure I: (left) Reconstructed image scaled by the intensity of the gradient with respect to different sections of the learned feature vector, showing the emergence of agent-aware feature representations. (right) In most RI. tasks, algorithm's 's Efficiency Normalized Score (see 4.3) improves when aug- mented with Ego-Foresight, sometimes outperforming supervised and model-based methods. The dashed line separates model-free from model-basee algorithms.
Although comparatively less explored, the idea of separately modeling the agent and the envi- ronment has also been investigated RL, with previous work demonstrating improved sample- efficiency in simulated control tasks (Gmelin et al., 2023). Additionally, this type of approach has been used to allow zero-shot policy transfer between different robots (Hu et al., 2022) and to im- prove environment exploration (Mendonca al., 2023). Common to all these works the reliance on supervision to obtain information about the appearance of the robot, allowing explicit disentan- gling of the agent from the environment. This usually provided in the form of mask of the agent within the scene, which is obtained either from geometric IDs in simulation, by fine-tuning segmentation model, or even by resorting to the CAD model of the robot. In a real-world robotic setting, it usually means the addition of separate system to segment the robot, adding complexity to the setup. Furthermore, supervised approaches are tied to the body-schema the agent and cannot adapt when it changes, for example, when using tools.
As humans, we not receive such detailed and hard-to-obtain supervisory signal and yet, during development, we build representation of ourselves (Watson, 1966) capable of adapting both slowly, as we grow, and fast, when we pick up tools (Maravita & Iriki, 2004). In this work, we argue that unsupervised awareness of self also be achieved in artificial systems, and study its advantages relative to supervised methods.
In approach, which name Ego-Foresight (EF), we place the agent's embodiment trinsic part of the learning process, since determines the visuomotor sensations that the agent can expect as it moves. Our insight is that agent information can be disentangled by having the agent move, while trying to predict the visual changes to its body configuration (Fig.l (left)), and that awareness and prediction of the agent's movements should improve its ability to solve complex tasks (Fig. I (right)).
in-
In our implementation, context RGB frames. A recurrent model takes the visual features and sequence of planned actions and predicts the future configurations of the agent, which the decoder reconstructs to obtain the future frames. This framework naturally lends itself to self-supervise training. In experiments with simulated data, we demonstrate the ability to predict the movement of the robot arm while ignoring moving objects in the environment and the integration of tools part of the body-schema.
convolutional encoder that receives input a limited amount of
Furthermore, we extend existing model-free and model-based RL algorithms with Ego-Foresighh, and assess the modified algorithms in two simulated domains and different control tasks involving a diverse range of embodiments, demonstrating that our approach can improve performance. We consider that two main factors contribute to this result: i) the disentanglement of agent information allows the RL algorithm to focus its capacity on learning the control the agent in the initial stages of training, and later on the external aspects and potential interactions within the
and
2

---

Published as a conference paper at ICLR 2026
environment, and ii) imposing predictability in the robot's movements regularizes the RL algorithm. In summary, the key contributions this paper are the following:
We propose self-supervised method for disentangling agent information based motion and self prediceeii,,on, removing the supervision required by previous methods.
We demonstrate the ability of our method to reconstruct future configurations of agent and adapt to changes in its body-schema.
We extend two state-of-the-ard RL algorithms (DrQ-v2 and TD-MPC2) with Ego-Foresight, showing improvement in sample-efficieng and performance in simulated control tasks.
We conduct an ablation study on the hyperparameters introduced by our approach.
BACKGROUND
Model-free visual RL RL problems are typically formulated as a Markov Decision Problem, de- fined as a tuple (S,A, T, R. 7, do), where S is the state space, .A is the action space, T(8++1 1 ar) the transition function, R(s,, a,) the reward function, y [o, 1] is a discount factor and dg is the distribution over the initial states Sy. The objective in RL is to learn the policy π :S - A that maximizes the expected discounted cumulative reward E, Cr 7 /R(ss,, a,)]. with a, (-1ss, and 82+1 ~ (-|8s,, az).
Over the last decade, work in RL from images (Mnih et al., 2013), where environment representa- tions are learned from high-dimensional inputs, has allowed agents to solve problems for which fea- tures cannot be designed by experts. widely used algorithm for RL from pixels is DDPG (Lillicrap et al., 2016), off-policy actor-critic algorithm for continuous control, which alternates between learning an approximator to the Q-function Qa and a deterministic policy тy. DrQ-v2 (Yarats et al., 2021) introduces augmentation technique for the input images, extending DDPG (Lillierap al., 2016) for improved efficiency. Additionally, it adopts the Twin Delayed variation of DDPG (Fu- jimoto al., 2018), which adds clipped double Q-learning and delayed policy updates to limit overestimation of the Q-values. Having sampled a batch of transitions T = (st. Gg 11 814n) from the replay buffer D, Q4 is learned by minimizing the mean-squared Bellman error:
Lerstie (b, D) E [(Qou (st, a;) - y) 4k E {1,2).
(1)
using target networks Qg oo slowly updated copies of the parameters Φx
to approximate the target values, with n-step returns, and where &R
n-1
Tree nd min Qi, (st-n, To (st4n)) k=1.2
The policy learned by maximizing EE~D (Qo To (s:)). to find the action that maximizes the Q-function. Improving sample-efficieng in RL from pixels remains open problem, and new methods for mitigating it have been consistently proposed (Hessel et al., 2018; Laskin et al., 2020; Stooke al., 2021; Zheng al., 2023).
Model-based visual RL In model-based RI., approaches such as Dreamer-v3 (Hafner et al., 2025) and TD-MPC2 (Hansen al., 2024) have achieved remarkable results by exploring the idea of training in an imagined latent space where thousands of parallel trajectories can be sampled with learned world model. In this work, we modify TD-MPC2 include Ego-Foresight as auxiliary loss term. The architecture of TD-MPC2 consists of five components:
Encoder: ht = E(st)
Policy prior: ā = π (h;)
Rew.: Fz = --[R(,,, az)
(3)
Dynamics:
h++1 = d(h;,
Term. value: q. Qh+, az)
Model updates done in two steps, with world-model updates using data from interactions stored in the replay buffer and optimizing the loss:
H X - sg (8.4 )) CE (frrr)) + CE (ât, (+))
CTDMPO(0,,)
=
E ToiN
(4)
t=o
"Media available https::/mserrranures thub. foresighr
3

---

Published as a conference paper at ICLR 2026
where 0 are the parameters of E, d, R and Q, À is temporally dependent discount factor, H is the MPC horizon, sg is the nop-gradient operator, CE is the cross-entropy loss, and qi is the TD-target. The policy prior is used to warm-start the planning procedure and accelerate its convergence and is optimized with maximum entropy. For planning. TD-MPPC uses Model Predictive Path Integral by taking advantage of the dynamics model rollout the latent trajectories that result from sampled action sequences and obtaining an estimate of their return using Q and R.
3
APPROACH
3.1 EPISODE PARTITION
We define our dataset as a collection of N episodes e, fixed length L (Fig. 2), which include sequence of observations in the form of RGB frames ir and the corresponding actions a of the agent:
e; (ZEg Bo..... (Iz, at)):.
1,..., N. During training, we randomly select a window of size C+ H within each episode which corresponds to the number of context frames plus the pre- diction horizon. This artificially aug- ments the available data by creating dif- ferent observations within each episode. Finally, from each window, C context steps are taken for input the model, as well as the actions for the whole se- quence up to tC+H. The frames be- tween to and fc+# are used as target for the prediction.
Xo.
Ny.
44
Xy. dg
Nq. ag
Xy.
Xg. ag 4
random window
iepursttions. Cisputrranes
Iupur actions Garget Crames
Figure 2:
Partition of an episode e,, From each se- quence of frames and actions (top), window of size C H is randomly sampled (middle). The first C steps of the window are used as context. For the remaining steps, the actions are used as input, while the frames serve as the target for the prediction.
3.2 MOTION- MOSION-BASED AGENT DISENTANGLEMENT
To model future visual configurations of the agent, we propose an encoder-decoder model (Fig. 3) with recurrent block that predicts the next agent configuration in feature space. The encoder, parameterized by it, produces a feature representation h E R", obtained from the context frames, h F.g (oit)..)) This vector is then split into scene h Rm and agent h² TR' features, with / + m = n. The agent features are used as input to the recurrent block, along with the action the next time step a for prediction of the next agent configuration. The recurrent network keeps predicting agent features 2 FCs (h' attt1 until time step tg. randomly sampled inside the prediction horizon in each training update. Finally, h is concatenated with h and passed the decoder for reconstruction of a De (htg h(te The components of the feature space corresponding h hold information about the visual content of the whole scene, which isn't present in h but is necessary for reconstruction. The components corresponding to h are then responsible for providing information to the decoder about the future configuration of the agent. Because scene information fast-forwrdded from te, the reconstructed frames tend to be close to r' except for the configuration of the agent for which there is a prediction at tg. This formulation results in the following reconstruction loss term:
Lef (vb D) E. T-D
[| x²
Crucially, we set the dimensionality leads the recurrent model to focus its capacity on the most predictable dynamics of the environment, which are those of the agent itself. Furthermore, by fast-forwarding the scene's 's content features h, encoded from the context frames the reconstruction time step tg. the recurrent block is discour- aged from predicting the dynamics of the complete environment, forcing the agent features to be concentrated on ha-
ht, to a small fraction of n, to create a bottleneck that
4

---

Published as a conference paper at ICLR 2026
h
te
h
tc
D
hfc+1
Ra
lk
s
cc
E
(...)
FC
foito
a c+11
fk-2
fk-1
Figure 3: Visuomotor prediction using Ego-Foresight.
3.3 AGENT VISUOMOTOR PREDICTION AS FEATURE LEARNING FOR RL
To test how method affects sample-efficieney in RI., implement it on top of two of the best performing RL algorithms for visual control: DrQ-v2 (Yarats et al., 2022), an off-policy model-free RL algorithm, and TD-MPC2, a model-based algorithm. We name the extended algorithms DrQv2- EF and TD-MPC2-EF, The episodes stored in the replay buffer D allow us to maintain the approach described in section 3.2 while jointly learning the policy. Despite our option for these algorithms, we preserve the generality of the approach so that it can be applied to any model-free or model- algorithm that makes use of experience replay and a visual encoder.
To extend an RL algorithm with Ego-Foresight, we take the feature vector at the output of the visual encoder and induce the disentanglement of agent information using the approach described in Section 3.2. Whereas the encoder is shared with the baseline algorithm, the recurrent unit and the decoder represent additional modules. An optimization step is then performed by augmenting the loss function of the baseline algorithm with Lef (5), with the losses backpropagating to the encoder.
Although this approach will fit most RL algorithms using a visual encoder, some may require more nuanced integration, for example, actor-critic algorithms in which the optimization is done in two steps. In the case of DrQ-v2 and TD-MPC2, the policy neural network is separate from the feature learning process, simply receiving the low-dimensional feature vector h4, coming from the encoder in the forward pass, while its gradients do not flow to the encoder in the backward pass. As such, in DrQ-v2 optimize Cef together with the critic loss of equation resulting in final objective function that is the weighted sum of both terms:
L(o, w. D) = (oritic D) BLe
(6)
In the original DrQ-v2 objective, the visual encoder is optimized via gradients coming from the critic network, and therefore learns to extract features that are predictive of value. The addition of Lef can be viewed as regularizer on the learned features, as it imposes them to be predictive of the agent's motion and is independent from the reward and the task.
In line with the model-free case, for TD-MPC2 we introduce Lef as an additional loss term in the model objective of equation 4, jointly optimizing
L(,, 1b D) LTDMPC(O, + BL. D).
When training Ego-Foresight together with an RL algorithm, we notice that the agent tends to per- form goal-directed movements as it seeks to maximize reward, preventing the observation of diverse enough motions to learn the visuomotor mapping. To solve this issue we introduce an optional motor-babbling (Saegusa et al., 2009; Kase et al., 2021) stage for fixed number of steps at the start of training, during which actions are random choices of +1, forcing exploratory movements. This means that the predictive ability for agent features is developed mainly during the initial stages of training. After the babbling stage, Ego-Foresight continues to be optimized according to equation 6, allowing adaptation in later stages, namely when the agent begins to pick-up available tool.
5

---

Published as a conference paper at ICLR 2026
Door Open
G.T.
Prediction
Hammer
G.T.
Prediction
12
16
20
24
1-28
Figure 4: The policy's solution to two tasks and the model prediction for the sequence of actions that solves the task. Notably, the hammer is integrated into the agent and therefore predicted.
4 EXPERIMENTS & RESULTS
4.1 EXPERIMENTAL SETUP
We conduct our experiments in simulation, with tests on two challenging RL benchmarks: Meta- world (Yu et al., 2020) and DMC (Tassa et al., 2018). The chosen benchmarks are implemented in MuJoCo (Todorov et al., 2012) and provide wide variety of tasks, ranging from robotic object manipulation to locomotion, and include a diverse set of embodiments. In all environments, we provide camera observations and the sequence of commanded actions inputs the model. In DMC we choose the 10 tasks for which official TD-MPC2 results from pixels are available online. For Meta- Woold, the tasks were chosen so that half would require the use of tools and the other half would require interaction with different kinds of objects.
4.2 AGENT- CHENT-ENVIRONMENT DISENTANGLEMENT
The experiments in this section intend to answer two fundamental questions in our study: (1) Is Ego- Foresight able to disentangle agent information and make predictions consistent with the ground truth? (2) Is our approach able adapt to changing embodiments?
Visualizing encoded features To study whether the proposed method for disentangling learns to concentrate information related to the agent in the feature vector ha, start by visualizing the areas of the reconstructed frames that are influenced the most by changes in the components of hg To do this, use DrQy2-EF and split the reconstructed frame into patches of 4 x 4 pixels and compute the gradient of each patch with respect to the section of interest of the feature vector at different moments during training run. We then scale the intensity of each patch proportionally to the gradient. In Fig. (left) it is possible to observe that at the start of training (after 10k steps), both h and h2 have similar influence in the predicted frame, with solid and static regions such as the background being quickly overfitted by the decoder. As training progresses, h, continues encode all the varying visual aspects of the scene, including the cabinet and the table borders (which change due to image augmentation), while h2 becomes specialized in agent information. We provide additional visualizations Appendix A.5.
In Fig. 4, we highlight two tasks using DrQv2-EF, showing the sequence of actions taken by the agent to solve them and how, for those action sequences, the motion the agent is correctly pre- dicted by EF. In the Door Open task, the robot starts by moving the gripper away from the camera toward the handle of the door and, after making contact with it, pulls back while opening the door. We highlight that in the predicted frames the door is reconstructed using visual content informa- tion from hts and therefore remains mostly static, as part of the scene. This contrasts with the Hammer task, where in addition to the movement of the agent, the model also predicts how the

---

Published as a conference paper at ICLR 2026
Assembby
Bankedball
Binn Picking
Box Close
Mean Eipisode Soccess
0.8
0.6
0.4
0.2
0.0
800% 0.0
0.0
800k
Hammer
Peg. Insert
Place
Stick
Moan Egisode Success
1.0
0.8
0.6
Mans
0.4
0.2
Environment Steps
Envirunment Staps
IM 0.0
Environment Steps
Environment Steps
1M
DRQV2
SEAR
CURL
DreamerV3
DRQv2-EF
Figure 5: Per-task results on WeradWorld tasks that require tool use.
hammer moves. This is due the fact that once the hammer is picked up, its motion is completely determined by the movement of the robot - effectively being integrated into the body of the robot and can therefore be predicted from the future actions. The same not true in the Door Open task, where there is no consistency between the movement of the door and that of the agent from episode to episode. This highlights the ability of our approach to adapt to changing embodiments, something that would not be possible in supervised setting. We note that in both visualizations, EF was trained with an horizon of 10 frames and used for prediction until time step 28, resulting in some divergence from the ground-trus -ttuh motion and in blurry predictions.
4.3 AGENT-AWARE REPRESENTATIONS FOR IMPROVED REINFORCEMENT LEARNING
In this section, we focus on the effect of Ego-Foresight on RL control tasks, intending to answer two central questions: (1) Can EF improve the sample-efficiency and performance of RL algorithms, in particular in tasks requiring tools? (2) How do the design choices affect the success of RL tasks?
Comparison Axes We present both per-task results and aggregate results across benchmarks and focus on sample-efficiency, measured in the number of environment steps and on performance in terms of success rate or reward, depending on the benchmark. For each task, we perform 5 runs per baseline using different random seeds and report mean and standard deviation (for TD-MPC2 on DMC we use the 3 seeds available in the official code repository). To measure sample-effiiiency and performance together, we define an Efficiency Normalized Score (ENS), with results presented in Fig. 1 (right). To compute ENS, we find the step at which 95% of the maximum performance achieved by any baseline on given task reached. Then, we measure the performance of each algorithm at that point. For robustness, this procedure is repeated for the 90% and 85% thresholds, with values being averaged across both thresholds and tasks obtain the final aggregate score.
Baselines With EF being designed to serve as an auxiliary task for improved feature learning in existing RL methods, the main baseline for an extended algorithm is the original algorithm itself. As such, we focus mainly on the comparison of DrQv2-EF with DrQ-v2 and TD-MPC2-EF with TD-MPC2. We provide three additional baselines: SEAR (Gmelin et al., 2023), CURL. (Laskin et al., 2020) and Dreamer-v3 (Hafner et al., 2025), introduced in Section SEAR builds on top of DrQ-v2, by splitting the feature representation into agent and full environment information, which achieved with a supervisory mask. SEAR is therefore a close supervised baseline to our approach. In the experiments below, the supervisory mask is obtained directly from the simulator. CURL uses a contrastive objective for self--suer visual feature extraction. Dreamer-v3 is a model-based RL. algorithm that has achieved state-of-the-art results in wide range of RL tasks. All results

---

Published as a conference paper at ICLR 2026
Acrobot Swingup
Cheetah Run
Finger
Finger Turn Easy
Mean Episods Reaand
1000
800
600
400
200
Finger Turn Hard
Quadruped Walk
Reacher Easy
Reacher Hard
Moan Episode Reward
1000
800
600
400
200-
Walker Run
Walker Walk
Mean
0.0 Environment Steps
IM
Mean Egioode Revard
1000
800
600
400
200
0.0
Environment Steps
1M 0.0
Environment Steps
0.)
Environment Steps
DreamerV3
TD-M
TD-MPC2-EF
Figure 6: Per-task results on DMC for the model- based RL algorithms.
obtained using the code provided by the authors of each paper. Unless otherwise noticed, we use the default hyperparameters of each baseline, including for DrQv2-EF and TD-MPC2-EF.
Reinforcement Learning Results In Figures 1, 5, 6 and in Appendix A.3 we present the aggregate and per-task results on 26 RL tasks. Additionally, we report Rliable metries (Agarwal et al., 2021) in Fig. 7 for rigorous evaluation of the statistical uncertainty of each baseline. We observe that in 21 of those tasks, augmenting DrQ-v2 with EF leads to improvements over the original algorithm, often reducing the amount of steps necessary to solve the task by significant margin and in many cases improving the asymptotic performance. Importantly, augmenting DrQ-v2 with EF does not degrade the results in any of the tasks. DrQv2-EF also shows improvement in the ENS and Rliable metrics over the supervised bascline SEAR and over CURL in both Meta- World and DMC and is competitive with Dreamer-v3, one of the best performing models in the literature. Still, in some tasks, these baselines outperform DrQv2-EF, as can be observed in Fig. 5. We attribute this to the diversity of reward functions of the benchmark, which means that some algorithms will adapt better to some tasks than to others. Interestingly, in tasks that require the use of tools (Fig. 5), we notice that the overall performance gap to the baselines increases, when compared with tasks that do not require the use of tools (see Fig. 1 and Appendix A.3). We interpret this result as a consequence of the tools being integrated in the feature representation of the agent, as their movement becomes predictable from the robot's actions when the agent is holding them.
Similar improvements are obtained for TD-MPC2-ER EF, which improves on out of tasks versus the original TD-MPC2, while obtaining similar results to the original algorithm in the other 2 tasks. In terms of ENS (Fig. 1), a significant improvement is achieved with the addition of our approach. Moreover, the results of TD-MPC2-EF indicate that our approach is not specific to single archi- tecture or to model-free RL, and that can benefit other algorithms in the visual RL community.
Ablation Study To better understand the effect of the different design choices in this work, we ablate the main hyperparameters introduced by Ego-Foresight on DrQv2-EF and present the results in Fig. 8. By varying the weight 3 of the EF loss term, is possible to observe how the proposed auxiliary loss has regularizing effect learning, improving performance without explicitly opti- mizing for reward. We also observe that the horizon length does not have a consistent impact on the results and that even for short horizons results improve DrQ-v2 (see Fig. 11). This indicates
8

---

Published as a conference paper at ICLR 2026
Meta-Worke World Rliable Metries
Median
JQM
Mea
DrQv2-EF
DrQqv2
SEAR
Dreamer-v3
0.64
0.72
0.80
0.88
0.72
0,80
0.80
0.6
0.7
Aggregate Success
DMC Rliable Metrics
Median
DiQuZ-EIF DrQ-v2
Dreamer-v3
ID-MPCZ
ID-MPC2-EF
720
800
730
840
900
700
750
Aggregate Return
Figure 7: Rliable metrics on the 16 Meta- World and 10 DMC tasks with 95% stratified bootstrap confidence intervals using the best evaluation score.
5000
Cor Weight
Prediction Horizon
Agent Feature Dim
Babbling Steps
Mean Erpisnds Reward
3000
9-00
dm Purtt
2000
de hunt
Super 10306
dim/hu644
Sagerd2300
1000
dim/uu122
Rob. Sopen13
B-40
din/ur5112
Bab Sagaro
0.0
Steps
500K 0.0 Favironment Steps
500K 0.0
Enviromment Steps
500K 0.0
Steps 500K
Figure 8: Ablation of the hyperparameters introduced in DrQv2-EF on Meta-World Door Open.
that the model retains the ability of learning the visuomotor mapping for smaller H. Nevertheless, horizon values of 10 and 40 show stronger performance, with H 10 being the preferred choice for default value due to lower computational cost. In terms of the dimensionality of the agent features h,, the ablation study shows that larger sizes have a detrimental effect, possibly because a softer bottleneck leads the recurrent block try predict other environment features, reducing the ability to disentangle agent information. Additionally, because the dimensionality of h is fixed at n and distributed between h2 and ha, going over the dimension size necessary to represent the state-space of the agent adds no benefit to agent representation while removing representational capacity from hg. Similarly, the introduction of babbling is meaningful contribution the results, but when applied for too many steps it can delay learning.
RELATED WORK
The notion of distinguishing self-generated sensations from those caused by external factors has been studied and referred under different terms and motivated by broad range of applications. Originating in psychology and neuroscience, with the study of contingency awareness (Watson, 1966) and sensorimotor learning (motor prediction) (Wolpert al., 2011; Wolpert & Flanagan, 2001), recently this concept has seen growing interest in Al as an auxiliary learning mechanism.
In developmental robotics, Zhang Nagai (2018) have focused on this problem from the standpoint of self-other distinction, by employing 8 NAO robots observing each other executing a set of motion primitives, and trying to differentiate self from the learned representations. Lanillos et al. (2020) note that to answer the question "Is this my body?", agent should first learn to answer "Am I generating those effects in the world?". Their robot learns the expected changes in the visual field it moves in front of a mirror or in front of twin robot and classifies whether it is looking at itself or not. Our approach is somewhat analogous to these works, the sense that we identify as being part of the agent what be visually predicted from the future actions while the robot moves.

---

Published as a conference paper at ICLR 2026
Still in robotics, the idea of modeling the agent has connections with work in body perception and visual imagination for goal-driven behaviour (Sancaktar et al., 2020). Another application that has been explored is the learning of modular dynamic models, that decouple robot dynamics from world dynamics, allowing the latter to be reused between robots with different morphologies. Hu et al. (2022) propose method for zero-shot policy transfer between robots, by taking advantage of robot- specific information such as the CAD model to obtain a robot mask from which future states can be predicted, given its dynamics. These future states are then used to solve manipulation tasks using model- based RL. Finally, this concept has been used to ignore changes in the robot as a means of measuring environment change and thereby incentivizing exploration (Mendonca et al., 2023).
In machine learning (ML), the distinction between the agent and the environment has been studied under the umbrella of disentangled representations, a long- -taaddinn problem in MI. (Bengio, 2013; Wang et al., 2024). While most works take an information theoretic approach to disentangled repre- sentation learning (Higgins et al., 2017; Kim & Mnih, 2018), some try to take advantage of known structural biases in the data, which is particularly relevant for sequential data, as it usually contains both time-variant and invariant features (Wiskott & Sejnowski, 2002). In video, this allows to dis- entangle content from motion (Villegas et al., 2017). Denton & Birodkar (2017) explore the insight that some factors are mostly constant throughout a video, while others remain consistent between videos but can change over time, to disentangle content and pose.
Finally, agent-environment disentanglement has seen growing interest in RI., being achieved through attention mechanisms (Choi al., 2018) or, more commonly, from explicit supervision (Gmelin et al., 2023; Qian et al., 2025; Kim et al., 2025). In Gmelin et al. (2023), the authors demonstrate that learning this distinction allows RL algorithms achieve better sample-efficiency and performance, serving as our most direct baseline. Vision transformers have also been adopted in the field, as they learn a feature space in which latent dynamics model can then be trained (Seo et al., 2023). Similarly to our work, some other approaches have combined ideas from both model-free and model- auxiliary losses (Racanière
based RI., by augmenting model-free algorithms with prediction- et al., 2017; Schwarzer et al., 2021). However, these works require learning a full dynamics model and do not consider any distinction between agent and environment.
6 CONCLUSIONS
Analysis and Limitations In this work, we studied how motion can be used to disentangle agent information a self supervisd manner. We integrate our approach with two RL algorithms, demon- strating improved sample-efficiency and performance, particularly in tasks requiring the use of tools. Furthermore, by removing the need for supervision while improving sample-efficiency, our method can be seen as step towards RL in real-world settings. We keep the generality of our approach, so that it can be used by other researchers on top of other RL algorithms. Some limitations of our work include the need for pixel-wise reconstruction of the scene, an auxiliary task that still generates debate in the field as to whether it imposes the need to model task-irrelevant or impossible to predict information (Hansen et al., 2022; LeCun, 2022). This could be addressed in the future with the use of a contrastive loss, which is computed in feature space and removes the need for reconstruction. The fixed babbling stage is another limitation, as in some tasks the reward immediately shoots up after babbling. In the future, this stage should be made adaptive. Furthermore, our approach still suffers from the characteristic instability of baseline RL algorithms such as DDPG (Islam et al., 2017: Henderson et al., 2018), an issue that hampers development and limits progress in the field. Finally, there is wall-clock overhead associated with EF (see Appendix A.4). Nevertheless, this penalty is compensated by the gains in sample-efficiency, which in applications like robotics is of greater importance than computational speed.
Future Work Future work directions include further exploration of the adaptation to new body- schemas, in particular how fast the model can adapt to changes and whether it can generalize to previously unseen tools. We also intend to take advantage of the disentangled information to improve modeling of moving objects in the scene. Another potential research direction is to scale experiments to higher resolution images of real-world environments, exploring ways to integrate Transformer architectures into feature learning for RL. Finally, our approach could also be applied to domains outside robotics, such as autonomous driving, where the main difference is that the agent's 's actions control the optical flow of the observed world and not the agent's body.
10

---

Published as a conference paper at ICLR 2026
ACKNOWLEDGMENTS
This study is in part funded by Lisbon ELLIS Unit, the Center for Responsible AI (PRR), LARSYS FCT funding (DOI: 10.540991A///0002222222 0.55499/UIDD/55000/2000
10.5499//UDD//50099222200 and YD acknowledges the support of UKRI grant EP/Y028732/1, and his RAEng CiET. This work was also supported by grants from NVIDIA and utilized NVIDIA RTX A6000 Ada graphics cards.
REFERENCES
Rishabh Agarwal, Max Schwarzer, Pablo Samuel Castro, Aaron C Courville, and Mare Bellemare. Deep reinforcement learning at the edge of the statistical precipice. Proc. of the Conf. on Neural Information Processing Systems (NIPS), 34:29304 29320, 2021.
Yoshua Bengio. Deep learning of representations: Looking forward. International conference on statistical language and speech processing, pp. 1-37. Springer, 2013.
Sarah-J Blakemore, Chris D Frith, and Daniel M Wolpert. Spatio-temporal prediction modulates the perception of self-produced stimuli. Journal of cognitive neuroscience, 11(5):551 1555,, 1999.
Sarah-Jayne Blakemore, Daniel Wolpert, and Chris Frith. Why can't you tickle yourself? Neurore- port, (1(1))RR1-RR 2000.
Jongwook Choi, Yijie Guo, Marcin Moczulski, Junhyuk Oh, Neal Wu, Mohammad Norouzi, and Honglak Lee. Contingency-awarc exploration in reinforcement learning. arXiv preprint arXiv: 18100448, 2018.
Andy Clark. Surfing Uncertainty: Prediction, Action, and the Embodied Mind. Oxford University Press, 2016. ISBN 9780190217013.
Remi Denton and Birodkar. Unsupervised learning of disentangled representations from video. In Proc. of the Conf. on Neural Information Processing Systems (NIPS), pp. 4414-4423, 2017.
Frederik Ebert, Chelsea Finn, Alex X Lee, and Sergey Levine. Self-supervised visual planning with temporal skip connections. Conference on Robot Learning (CoRL.), 2017.
Chelsea Finn, Ian Goodfellow, and Sergey Levine. Unsupervised learning for physical interaction through video prediction. Proc. of the Conf. on Neural Information Processing Systems (NIPS), 29, 2016.
Karl Friston. The free-energy principle: a unified brain theory? Nature reviews neuroscience, (2):177-18,, 2010.
Scott Fujimoto, Herke Hoof, and David Meger. Addressing function approximation error in actor- critic methods. In International Conference on Machine Learning (ICML), pp. 1587-1596.
PMLR, 2018.
Kevin Gmelin, Shikhar Bahl, Russell Mendonca, and Deepak Pathak. Efficient RL via disentan- gled environment and agent representations. In International Conference on Machine Learning (ICML), volume 202 of Proceedings of Machine Learning Research, pp. 11525-11545. PMLR, 2023.
Christian Gumbsch, Noor Sajid, Georg Martius, and Martin V Butz. Learning hierarchical world models with adaptive temporal abstractions from discrete latent dynamics. In International Con- ference on Learning Representations (ICLR), 2024.
David Ha and Jürgen Schmidhuber. Recurrent world models facilitate policy evolution. In Proc. of the Conf. on Neural Information Processing Systems (NIPS), 2018.
Danijar Hafner, Jurgis Pasukonis, Jimmy Ba, and Timothy Lillicrap. Mastering diverse control tasks through world models. Nature, 640(8055):647-633. 2025.
Nicklas Hansen, Xiaolong Wang, and Hao Su. Temporal difference learning for model predictive control. arXiv preprint arXiv:2203.0490 044955 2022.

---

Published as a conference paper at ICLR 2026
Nicklas Hansen, Hao Su, and Xiaolong Wang. Td-mpc2: Scalable, robust world models for contin- wous control. In Interaational Conference on Learning Representations (ICLR), 2024.
Peter Henderson, Riashat Islam, Philip Bachman, Joelle Pineau, Doina Precup, and David Meger. Deep reinforcement learning that matters. In Proc. of the Conference on Advancements of Artifi- cial Intelligence (AAAI), volume 32, 2018.
Matteo Hessel, Joseph Modayil, Hado Van Hasselt, Tom Schaul, Georg Ostrovski, Will Dabney, Dan Horgan, Bilal Piot, Mohammad Azar, and David Silver. Rainbow: Combining improvements in deep reinforcement learning. Proc. of the Conference on Advancements of Artificiil Intelligence (AAAT), volume 32, 2018.
Irina Higgins, Loic Matthey, Arka Pal, Christopher P Burgess, Xavier Glorot, Matthew M Botvinick, Shakir Mohamed, and Alexander Lerchner. beta-vae: Learning basic visual concepts with con- strained variational framework. International Conference on Learning Representations (ICLR), 3, 2017.
Edward S. Hu, Kun Huang, Oleh Rybkin, and Dinesh Jayaraman. Know thyself: Transferable visual control policies through robot-awareness. International Conference on Machine Learning (ICML), 2022.
Riashat Islam, Peter Henderson, Maziar Gomrokchi, and Doina Precup. Reproducibility of benchmarked deep reinforcement learning tasks for continuous control. arXiv preprint 10t.:1778.043333 2017.
Kei Kase, Noboru Matsumoto, and Tetsuya Ogata. Leveraging motor babbling for efficient robot learning. Journal of Robotics and Mechatronics, 33(5):1033-1074, 2021.
Hyunjik Kim and Andriy Mnih. Disentangling by factorising. In International Conference on Machine Learning (ICML), pp. 2649-2658. PMLR, 2018.
Kyungmin Kim, JB Lanier, Pierre Baldi, Charless Fowlkes, and Roy Fox. Make the pertinent salient: Task-relevant reconstruction for visual control with distractions, 2025.
Pablo Lanillos, Jordi Pagès, and Gordon Cheng. Robot self/other distinction: active inference meets neural networks learning mirror. In European Conference on Artificial Intelligence, 2020.
Michael Laskin, Aravind Srinivas, and Pieter Abbeel. Curl: Contrastive unsupervised representa- tions for reinforcement learning. International Conference on Machine Learning (ICML), pp- 5639-5650. PMLR, 2020.
Yann LeCun. A path towards autonomous machine intelligence version 0.9. 2, 2022-06-27. Open Review, 62(1):1-62, 2022.
Timothy Lillicrap, Jonathan J. Hunt, Alexander Pritzel, Nicolas Heess, Tom Erez, Yuval Tassa, David Silver, and Daan Wierstra. Continuous control with deep reinforcement learning. In Inter- national Conference on Learning Representations (ICLR), 2016.
Angelo Maravita and Atsushi Iriki. Tools for the body (schema). Trends in cognitive sciences. 8(2): 79-86, 2004.
Russell Mendonca, Shikhar Bahl, and Deepak Pathak. Alan: Autonomously exploring robotic agents in the real world. Proc. of the IEEE Intl. Conf. on Robotics Automation (ICRA), pp. 3044- 3050. IEEE, 2023.
Volodymyr Mnih, Koray Kavukeuoglu, David Silver, Alex Graves, Ioannis Antonoglou, Daan Wier- stra, and Martin Riedmiller. Playing atari with deep reinforcement learning. arXiv preprint 112.5233 2013.
Long Qian, Ziru Wang. Sizhe Wang, Lipeng Wan, Zeyang Liu, Xingyu Chen, and Xuguang Lan. Pre-training robo-centric world models for efficient visual control, 2025. URL https: / //openreview net
www.enntttooommmeeeeee forum? d
12

---

Published as a conference paper at ICLR 2026
Sébastien Racanière, Théophane Weber, David Reichert, Lars Buesing, Arthur Guez, Danilo Jimenez Rezende, Adrià Puigdomènech Badia, Oriol Vinyals, Nicolas Heess, Yujia Li, et al. prrggnnationaugmmntte agents for deep reinforcement learning. Proc. the Conf. Neural Information Processing Systems (NIPS),30 30, 2017.
Alec Radford, Luke Metz, and Soumith Chintala. Unsupervised representation learning with deep convolutional generative adversarial networks. International Conference on Learning Represen- tations (ICLR), 2016.
Olaf Ronneberger, Philipp Fischer, and Thomas Brox. U-net: Convolutional networks for biomed- ical image segmentation. In Medical image computing and computer-assiste aassisted intervention- MICCAI 2015: 18th international conference, Munich, Germany, October 5-9, 2015, proceed- ings, part 18. pp. 234-241. Springer, 2015.
Ryo Saegusa, Giorgio Metta, Giulio Sandini, and Sophie Sakka. Active motor babbling for senso- rimotor learning. In EEE International Conference on Robotics and Biomimetics, pp. 794-799. IEEE, 2009.
Cansu Sancaktar, Marcel AJ Van Gerven, and Pablo Lanillos. End-to-end pixel-based deep ac- tive inference for body perception and action. In International Conference Development and Learning and Epigenetic Robotics, pp. 1-8. IEEE, 2020.
Max Schwarzer, Ankesh Anand, Rishab Goel, Devon Hjelm, Aaron Courville, and Philip Bach- man. Data-efficient reinforcement learning with self-predictive representations. International Conference on Learning Representations (ICLR), 2021.
Younggyo Seo, Danijar Hafner, Hao Liu, Fangchen Liu, Stephen James, Kimin Lee, and Pieter Abbeel. Masked world models for visual control, In Conference on Robot Learning (CoRL), pp. 1332-1344. PMLR, 2023.
Adam Stooke, Kimin Lee, Pieter Abbeel, and Michael Laskin. Decoupling representation learning from reinforcement learning. In International Conference on Machine Learning (ICML), pp- 9870-9879. PMIR, 2021.
Yuval Tassa, Yotam Doron, Alistair Muldal, Tom Erez, Yazhe Li, Diego de Las Casas, David Bud- den, Abbas Abdolmaleki, Josh Merel, Andrew Lefrancq, et al. Deepmind control suite, arXiv preprint arXiv:1801.0006 2018
Emanuel Todorov, Tom Erez, and Yuval Tassa. Mujoco: A physies engine for model-based control. In Proc. of the JEEE/RSJ Int. Conf. on Intelligent Robots and Systems (IROS), pp. 5026-5033 IEEE, 2012.
Ruben Villegas, Jimei Yang, Seunghoon Hong, Xunyu Lin, and Honglak Lee. Decomposing tion and content for natural video sequence prediction. International Conference on Learning Representations (ICLR), 2017.
Xin Wang, Hong Chen, Zihao Wu, Wenwu Zhu, et al. Disentangled representation learning, IEEE Trans. on Pattern Analysis and Machine Intelligence (TPAMI), 2024.
Eiji Watanabe, Akiyoshi Kitaoka, Kiwako Sakamoto, Masaki Yasugi. and Kenta Tanaka. Illusory motion reproduced by deep neural networks trained for prediction. Frontiers Psychology, Volume 9 2018, 2018. ISSN 1664-1078.
John S Watson. The development and generalization of "contingency awareness" in early infancy: Some hypotheses. Merrill-Pulmer Quarterly of Behavior and Development, 12(2):123-133) 1966.
Laurenz Wiskott and Terrence Sejnowski. Slow feature analysis: Unsupervised learning of invari- ances. Neural computation, 14(4):715-770 2002.
Daniel M Wolpert and Randall Flanagan. Motor prediction. Current biology. 11(18):R729-R732 2001.
Daniel M Wolpert, Jöm Diedrichsen, and Randall Flanagan. Principles of sensorimotor learning.
Nature reviews neuroscience, (2 12(12):739-551 2011.
13

---

Published as a conference paper at ICLR 2026
Philipp Wu, Alejandro Escontrela, Danijar Hafner, Pieter Abbeel, and Ken Goldberg. Daydreamer: World models for physical robot learning. In Conference on Robot Learning (CoRL), pp. 2226- 2240. PMLR. 2023.
Denis Yarats, Ilya Kostrikov, and Rob Fergus. Image augmentation is all you need: Regularizing deep reinforcement learning from pixels. In International Conference on Learning Representa- tions (ICLR), 2021.
Denis Yarats, Rob Fergus, Alessandro Lazaric, and Lerrel Pinto. Mastering visual continuous con- trol: Improved data-augmented reinforcement learning. International Conference on Learning Representations (ICLR), 2022.
Tianhe Yu, Deirdre Quillen, Zhanpeng He, Ryan Julian, Karol Hausman, Chelsea Finn, and Sergey Levine. Meta-world: A benchmark and evaluation for multi-task and meta reinforcement learning. In Conference on Robot Learning (CoRL), pp. 1094-1100. PMLR, 2020.
Yihan Zhang and Yukie Nagai. Proprioceptive feedback plays key role in self-other differentiation. In International Conference Development and Learning and Epigenetic Robotics, pp. 133-138.
IEEE, 2018.
Ruijie Zheng, Xiyao Wang. Yanchao Sun, Shuang Ma, Jieyu Zhao, Huazhe Xu, Hal Daumé III, and Furong Huang. Taco: Temporal latent action-driven contrastive loss for visual reinforcement learning. Proc. of the Conf. Neural Information Processing Systems (NIPS), 36:48203 48225,, 2023.
14

---

Published as a conference paper at ICLR 2026
A APPENDIX
A.1 ARCHITECTURE
When extending an algorithm with EF, we preserve the architecture of the visual encoder, as this is shared between the original algorithm and EF. Because our approach requires the introduction of a bottleneck, in DrQv2-EF we introduce an additional average pool and linear projection at the output of the encoder, to reduce the dimensionality of the DrQ-v2 feature vector from 39200 to 2048. For TD-MPC2, the encoder architecture completely preserved. The recurrent unit is based on the forward dynamics model of TD-MPC2 which consists of an MLP with LayerNorm and Mish activations and has demonstrated strong results latent space prediction. The decoder is based on DCGAN (Radford al., 2016).
A.2 REAL-WORLD ENVIRONMENT
To test the real-world transferability of Ego-Foresight, we perform qualitative assessment the predictions on the BAIR dataset (Ebert et al., 2017), which consists of a Sawyer robot randomly pushing wide range of objects a table. In Fig. 9, we show two sample predictions from Ego- Foresight. The model succeeds in separating visual information that is part of the robot from the scene, predicting the trajectory of the robot's 's arm according to its motion, which shows that the model correctly learned the visuomotor map between actions and vision. The background, including moving objects, reconstructed their original position, as encoded in h Even so, Ego-Foresight still predicts that some change should happen when the arm passes by an object, and it often blurs objects that predictably would have moved. However, it should be noted that the agent actions and object movement are inherently correlated and therefore cannot be completely disentangled.
We also evaluate the ability to generalize to previously unseen trajectories. To achieve this, we handcraft a previously unseen movement, as shown in Fig. 10. While we display single handcrafted example, this simple experiment shows that, provided a policy function, the dynamics learned by Ego-Foresight should manage to predict the outcome of sampled trajectories.
In experiments with the BAIR dataset, both the encoder and the decoder are based on DCGAN. This enables the addition of U-Net like skip-connections (Ronneberger et al., 2015), which facilitate the reconstruction of fine geaigraindd details and allow the prediction to focus on the agent. These improve the prediction quality without impacting performance and can optionally be used in the simulated experiments too.
G.T.
Pred.
G.T.
Pred.
1=
E13
F17
E21
1=25
1-29
Figure 9: Predictions the BAIR Dataset.
15

---

Published as a conference paper at ICLR 2026
EP
Context Frame
Original trujectory
Output
EF
Context Frame
Handerafted trajectory
Output
Figure 10: Generation an unseen trajectory.
A.3 SUPPLEMENTARY PER-TASK RESULTS
We present per-task results on the remaining Meta-World tasks that do not require tools in Fig. 11 and the results of the model-free algorithms on DMC described in Section 4.3 in Fig. 12.
Bation Press Topdowe Wall
Butson Press Well
Door Close
Dogr Lock
0.8
0.6
0.2
0.0
0.0
400k 0.0
500k 0.0
400% 0.0
500k
Unloock
Pull
Episode
0.6
D&
0.2
0.0
0.0
Enviromment Steps 500% 0.0
Steps 500k 0.0
Steps 500k
Environment Steps 800k
DRQv2
SEAR
CURL
DreamerV3
DRQv2-EF
Figure 11: Per-task results Meta-World tasks requiring tools.
16

---

Published as a conference paper at ICLR 2026
Acrobot Swingup
Cheetah Run
Finger Spin
Finger Turn Easy
Episode Reward
1000
500
600
400
Wy
200
Finger Tum Hard
Quadruped Walk
Reacher Easy
Reacher Hard
Episode Rewart
1000
800
mmM
600
400
200
Walker Run
Walker Walk
Mean
0.0
Environment Steps 1M
Egisode Reward
1000 1000
800
600
400
200
Enviromment Steps
Environment Steps
Environment Steps
SEAR
DRQV2
CURL
DRQv2-EF
Figure 12: Per-task results DMC for the model-free RL algorithms.
A.4 TRAINING DETAILS
We opt use the official results provided in tables by the authors of each baseline, whenever avail- able, in an effort to save resources. This was the case in the DMC benchmark for DrQ-v2, Dreamer- v3 and TD-MPC2. In all other instances, results obtained by executing the source code provided by the authors of each work. This is done using NVIDIA RTX A6000 and RTX 6000 Ada GPUs but cheaper hardware too, such
the models be necessary for extending an algorithm with Ego-Foresigge are also implemented on top of the original code, without changing the default hyperparameters. ters introduced by Ego-Foresighe is provided in Table 1. For the hyperparameters that were part of the original DrQ-v2 and TD-MPCK we use the default values provided in the respective papers and source code. Similarly, for SEAR and Dreamer-v3, all hyperparameters are the default ones.
NVIDIA RTX 4060. The modifications
list of the values used for the hyperparame-
In Table 2 we present an analysis of computational cost each algorithm terms of number of environment steps taken per second.
Table 1: Default hyperparameters used in DrQv2-EF and TD-MPC2-EF.
Ego-FFrresight Parameters
Parameter
DrQv2-EF
TD-MPC2-EF
Hidden dim.
2048
512
h, dim.
2016
480
h. dim
32
32
Motor-babbling steps
25000
1250
Prediction horizon (time steps)
10
10
17

---

Published as a conference paper at ICLR 2026
Table Environment steps per second obtained on one NVIDIA RTX A6000 using batch size 256 and action repeat 2 on the DeepMind Control benchmark, For Dreamer-v3, the default batch size of 16 was used.
Algorithm
Env. steps per second
DrQ-v2
155
DrQv2-EF
110
SEAR
CURL
44
Dreamer-v3
41
TD-MPC2
28
TD-MPC2-EF
23
A.5 ADDITIONAL VISUALIZATIONS
In Fig. 13 we provide additional visualizations the learned features on the Hammer task Meta- World. Feature visualizations are obtained with the method presented in Section 4.2. Each row of Fig. 13 presents a snapshot of the learned features at different training moments, in particular one at the beginning of training, another after the babbling stage and another after the agent has learned to pick up the hammer. From columns (3) and (5) it is possible to visualize how our approach learns an approximation
the true agent mask. Additionally, we note how the hammer begins to be rep- resented together with the robot once it is gripped. Finally, in column (4) we present a thresholded version of (4), allowing more direct comparison with (5). The pixelated appearance these frames is a product of the patching step involved in obtaining the visualizations (see Section 4.2). We com- pute precision and recall between columns (4) and (5) obtaining the pairs: 5k steps (0.11, 0.35), at 30k steps (0.32, 0.30), at 500k steps (0.18, 0.17). We do note that these values are affected by the size of the patches and by the thresholding value, and are merely indicative.
(2) Pred. scaled by (3) Pred. scaled by (4) Thresholded grad. w.r.l scene grad. w.r. agent visualization of 3. features ha
(5) Ground truth agent mask
(I) Ground Truth
features ht
At 5k train steps
At 30k train steps
At 700K train steps
Figure 13: Visualization of the learned features on the Hammer task of Meta- World.
18


Reviews:
[{"review_id": "6SLgvcfErJ", "reviewer": "Reviewer_eUBc", "summary": "This submission proposed Ego-Foresight (EF), a self-supervised method using motion prediction to disentangle agent information from environment, aiming at improving sample efficiency for deep RL. The method uses an encoder-decoder architecture with a recurrent module that predicts future agent configurations from actions, with a bottleneck forcing focus on agent-specific features. The authors integrate EF with DrQ-v2 and TD-MPC2, and compare with SEAR and Dreamer on Meta-World and DMC benchmarks, showing improved sample efficiency particularly on tasks requiring tool use.", "questions": "1. Can you explain more what Fig 4 is implying? Perhaps using the same visualization for Fig 1 to highlight what’s actually being learned by the agent feature?\n2. In the ablation study, larger agent feature dimension seems to have a detrimental effect, which is attributed to “reducing the ability to disentangle agent information.” Might an alternative explanation be just because when the total feature dimension is too large, the policy is harder to learn as the state space is now larger?\n3. For tasks where EF doesn't help (e.g., Walker tasks), what is the failure mode? Can you characterize when EF is expected to help?\n4. How good is the learned agent feature? Since EF is proposed to disentangle agent from the environment, can you quantify how much overlap there is between the learned agent feature with the supervised agent masks if available in some tasks?\n5. Tool embodiment is mentioned as an interesting capability, and is perhaps when EF might be helpful the most. Does this adaptation happen within a session? Say in the hammer task, does the agent feature starts to include the hammer only after the hammer is being picked up? Or does the agent feature already include the hammer from the beginning?", "limitations": "", "rating": 4, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "1. The proposed self-supervised framework for learning agent-aware representation, as opposed to supervised training that requires agent masks in previous work.\n2. The proposed method can be integrated with both model-free (DrQ-v2) and model-based (TD-MPC2) algorithms, showing flexibility.\n3. The evaluation is comprehensively done on 26 tasks across two benchmarks and baseline methods, including Dreamer (model-based), SEAR (previous supervised agent-aware learning method). Results are presented with proper statistical analysis with RLiable metrics, and the ablation studies are useful for dissecting the proposed model.\n4. The paper is generally well-written with good use of figures, including the model diagram.", "weaknesses": "1. The results on performance improvement does not seem to be consistent across the tasks. In figures 5 and 6, SEAR and Dreamer seem to perform equally well, and sometimes even better, than EF. This performance gap seems to be more prominent in figures 11 and 12. This inconsistency is not adequately explained. Also, for efficiency normalized score presented in figure 1, could you please provide standard deviation/ errors to properly assess the improvement? \n2. Since the proposed EF algorithm is a self-supervised algorithm for image based deep RL, why are other self-supervised image-based RL algorithms not included as comparison baselines, such as CURL (Srinivas, Laskin, et al, 2020) and SPR (Schwarzer, Anand, et al, 2021)? This and the performance inconsistency are perhaps the main weaknesses of the current draft.\n3. The ablation study is performed on the DoorOpen task, which is not part of the MetaWorld tasks that require tool use as presented in figure 5. Since you highlighted EF should be better in tasks that require tool use, why is the ablation study done on DoorOpen?\n4. As learning efficiency is emphasized, it’d be good to provide comparisons on computational overhead (wall-clock time, memory, etc.), which is mentioned but not quantified.\n5. Maybe I missed something, but in terms of the algorithm (neuroscience inspiration aside), Intuitively, why would learning disentangled representation lead to improved sample efficiency?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "RXjZZhIkyR", "reviewer": "Reviewer_Nv8q", "summary": "This work proposes idea of self-supervised learning of consequences of RL agent's actions in an environment for inducing agent-aware representations. The motivation behind the work is grounded in how humans foresee the effects of their actions while executing them which allows them to perform the tasks more effectively. The paper proposes use of convolutional encoder-decoder architecture that predicts future frames given the visual features and sequence of planned actions. A self-supervised loss is applied on the predictions that allows us to learn visual features that enable easy prediction of the future, and hence instill agent-awareness. The authors extend pre-existing RL algorithms in robotics with proposed ego-foresight loss that allows for improved performance.", "questions": "Are these results also applicable in non-robotics environments like standard control tasks (e.g. CartPole) in RL? It would be nice to have an auxiliary loss that allows the neural network to predict the next states to align the representations used for deciding the policy. It is only a minor suggestion.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "1. The paper is very well written with the introduction making connection of the technique with human locomotion and various other foresight instances. \n2. The approach proposed in this work seems intuitive and should be easy to integrate into existing RL pipelines learning from visual features.\n3. The results demonstrate consistent improvement when the technique is used in conjunction with prior frameworks. \n4. The choice of the research questions studied is apt for this work and the paper goes into detail answering each of them.\n\nOverall, I find the quality of the write-up, the experiments and the results high and recommend accepting the work.", "weaknesses": "1. The paper uses recurrent neural networks, while the overall efficacy of transformers in sequence generation is well-proven in the domains of natural language and vision.\n2. Episode partitioning uses a fixed constant H for defining the future window. This choice could have been made more dynamic to remove any bias due to burden of predicting next H frames every time.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "QtySfYh54a", "reviewer": "Reviewer_Y4MG", "summary": "This paper proposes Ego-Foresight (EF), a self-supervised method, where the core idea is to disentangle the agent’s representation from the environment’s representation by leveraging motor prediction. The model learns to separate a feature vector h into agent-specific features ha  and scene-specific features hs. A recurrent model predicts future agent features based on unrolled action sequences, while the scene features hs are fast-forwarded from the context frame. The predicted agent features and the original scene features are then used to reconstruct the predicted frame to train the encoder. This method avoids the need for explicit supervision required by prior work like SEAR. The authors demonstrate that adding EF as an auxiliary loss to model-free (DrQ-v2) and model-based (TD-MPC2) algorithms improves performance and sample efficiency on Meta-World and DMC tasks.", "questions": "## Questions\n1. Following Weakness #1: Could the authors please clarify the role of $h_s$? Is it intended to be static? Why does the gripper disappear in the Figure 4 prediction, and what does this imply about what $h_s$ is actually encoding? \n\n2. Following Weakness #2: Could the authors justify their claim that a ~500 reward gap in the horizon ablation is \"not significant\"? It appears to be one of the more significant factors in the study.\n\n3. The qualitative results in Figure 4 show predictions for relatively small movements. How does the model's predictive quality and disentanglement hold up when the agent performs large-scale actions or when predicting over a longer horizon?\n\n4. in Fig. 5 and 6, the authors do not integrate their approach into the Dreamer architecture, but only in the other two approaches. Why is that?", "limitations": "", "rating": 4, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "- The core idea of using motor-prediction as a self-supervised signal for agent-environment disentanglement is interesting. It provides an alternative to supervised methods that rely on segmentation masks, which are often unavailable in real-world settings.\n- The proposed EF module is model-agnostic and can be integrated as an auxiliary loss into both model-free and model-based RL algorithms, as demonstrated by its successful application to DrQ-v2 and TD-MPC2.\n- The method shows clear improvements in sample efficiency and asymptotic performance against strong baselines (DrQ-v2, TD-MPC2) and is competitive with its supervised counterpart (SEAR) on several tasks.", "weaknesses": "- Figure 4 (“Door Open” and \"Hammer\" task) seems to be central to the claims made in the paper and tries to show that the door is part of the environment whereas the hammer becomes part of the body schema. However, the distinction is not so clear to me. the hammer has to be picked up first, just like the agent has to grip the door handle. Once the contact is made the door could be moved to and fro just like the hammer. The illustrations in the figure also does not really help. It seems like the gripper disappears in the door opening task, i can't even see whether the door is open or close in the picture...\n\n- The authors state that “horizon length does not significantly impact the results”. However, the “Prediction Horizon” graph in Figure 8 clearly shows that H=10 and H=40 (peak reward ~4500) perform significantly better than H=5 and H=20 (peak reward ~4000). A ~500 reward difference (or ~12.5%) is not insignificant. It is similiar to the difference between β=0  (no EF module) and other β where it is claimed that the difference is important. This difference in treatment of the same kind of error difference seems concerning wrt to the paper's claims.\n\n Minor issues:\n\n- Presentation: The paper's structure could be tightened. Section 2 (\"Background\") is a single, long subsection that could be broken up. Similarly, for Section 4.1 (\"Experimental Setup\"), the special subtitle \"Environments\" is not necessary.\n\n- Typo: On Line 238, the feature vector is referred to as $h_{t_c}$. To be consistent with the paper's notation defined on Line 191, this should be $h^{t_c}$.\n\n- Clarity: In Figure 2, the \"action\" label is a bit vague. It seems strange to use a single variable $x$ to represent the observation, and use a vector $[a^0, \\cdots, a^M]$ to represent the action.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "dsmNlBN0Lk", "reviewer": "Reviewer_hc2R", "summary": "This paper proposes Ego-Foresight (EF), a training method that employs an auxiliary task of predicting agent-specific features. The authors show that EF improves sample efficiency when applied to both model-free and model-based RL algorithms. They show that EF lifts limitations of supervised methods for representation disentanglement and enables agents to adapt to changes in their body-schema, such as during tool use.\n\nThe authors provide strong motivation for their method and the results are convincing. The contribution of a self-supervised approach for motor prediction is of interest to the RL community. I have some concerns regarding the clarity of the mechanisms underlying EF's performance improvements, but overall, I think this paper should be accepted.", "questions": "**Questions for clarification:**\n\n1. The claim that \"predicting the sensory consequences of one's own movement is remarkably simpler, yet equally important\" cites Wolpert and Flanagan (2001), but this source doesn't appear to directly support this claim. Could you clarify what you mean by \"simpler\" and \"important\" in this context, and provide additional citations if available?\n2. You mention \"both h_a and h_s have a similar influence in the predicted frame, with solid and static regions such as the background being quickly overfitted by the decoder,\" and mention the \"regularizing effect on learning\". Could you elaborate on what constitutes overfitting in this context and how you observe the regularizing effect?\n3. You mention \"the need for pixel-wise reconstruction of the scene, which could be addressed in the future by exploiting a contrastive loss.\" Could you elaborate on this idea and how a contrastive loss might address the current limitations?\n4. You mention \"the characteristic instability of baseline RL algorithms such as DDPG.\" Did you observe this instability in your experiments?\n\n---\n\n**Minor comments:**\n\n- Figure 4 (left) could be clearer. It's difficult to distinguish the trajectory of the end effector because the end-effector is blurred. The trajectory was more distinguishable in the gif provided in the codebase, so maybe a different angle would be helpful for the static figure.\n- The claim that \"in terms of ENS, a significant improvement is achieved with the addition of our approach\" would be strengthened by adding confidence intervals to Figure 1 to support statistical significance.\n- In the ablation study you mention, \"In terms of the dimensionality of the agent features h_a, the ablation study shows that larger sizes have a detrimental effect.\" This wasn't immediately apparent to me because the larger dimensionalities appeared to have similar performance curves to the default value of 32.\n\n---\n\n**Suggestions for improvement (not factored into score):**\n\n1. Could evaluate EF with DreamerV3 to demonstrate generalization across model-based algorithms.\n2. Include additional analysis demonstrating the proposed distinct phases where the model learns agent-specific versus environment-specific features.\n3. Compare the performance boost from tool use versus no tool use, and compare this difference with the difference given by the supervised baseline (SEAR). Given your claims about supervised approaches being non-adaptable to body-schema changes, EF should demonstrate better performance gains on tool-use tasks compared to SEAR.\n4. The ablation study performance curves on the Door Open task show high variance with similar results across conditions. Could consider conducting ablations on tasks with larger differences, such as bin picking or peg insertion, to better distinguish the impact of different hyperparameters.\n5. You mention that sample efficiency is more important than computational efficiency, but it would still be helpful to quantify the wall-clock time tradeoff of EF.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "EF demonstrates improved sample efficiency when integrated with both model-free and model-based RL algorithms. The performance differences and improved sample efficiency between DRQv2-EF/TD-MPC-EF and their baseline counterparts are convincing across the tasks presented. Additionally, reporting Rliable metrics in addition to the performance curves add robustness to the results since they report statistic uncertainty more rigorously.\n\nEF successfully reconstructs future agent configurations and adapts to changes in body-schema such as tool use (Figure 4). This approach offers advantages over supervised methods, which are constrained by fixed body-schema and cannot adapt to changes such as tool use.", "weaknesses": "The authors propose that EF's improvements stem from (i) disentanglement of agent information allowing the RL algorithm to focus on agent control in early training stages and environmental interactions later, and (ii) regularization through imposing predictability in robot movements. However, there seems to be not enough evidence supporting this mechanism, particularly the claim about distinct training phases. While Figure 1 (left) shows some training progression, the predictive ability for agent-specific features appears to improve later in training rather than in the initial stages, which seems inconsistent with the proposed explanation.\n\nI also have a slight concern about the use of Efficiency Normalized Score (ENS). ENS identifies the step at which 95% of maximum performance is reached and measures algorithm performance at that point. This metric may inadvertently favor algorithms with higher performance variance which could reach the threshold faster than more stable learners. So ENS may not be reflecting true sample efficiency accurately.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)