Ego-Foresight: Self-supervised Learning of Agent-Aware Representations for Improved RL

1. Problem Statement

Deep Reinforcement Learning (RL) agents, particularly those operating from raw pixel inputs, suffer from poor sample efficiency, requiring vast amounts of training data to learn effective policies. A primary bottleneck is the difficulty in distinguishing between the agent's own dynamics (self-generated sensory changes) and environmental dynamics (external events).

• Current Limitations: Previous approaches that attempt to separate agent and environment information (e.g., SEAR) rely heavily on supervised signals, such as segmentation masks derived from geometric IDs, CAD models, or fine-tuned segmentation networks. These methods are brittle, difficult to deploy in real-world settings where such masks are unavailable, and cannot adapt when the agent's "body schema" changes (e.g., picking up a tool).
• Human Inspiration: Humans learn to distinguish self from environment through motor prediction (predicting the sensory consequences of one's own motor commands) without explicit supervision. This paper aims to replicate this capability in artificial systems to improve RL efficiency.

2. Methodology: Ego-Foresight (EF)

The authors propose Ego-Foresight (EF), a self-supervised learning framework that disentangles agent information from environmental information based on motion and prediction.

Core Architecture

EF utilizes an Encoder-Decoder architecture with a recurrent block:

1. Input: A sequence of context RGB frames ($x_{t_0:t_c}$) and a sequence of planned actions ($a_{t_c+1:t_k}$).
2. Feature Disentanglement: The visual encoder ($E_\psi$) maps context frames into a latent feature vector $h$. This vector is split into two components:
   • Scene Features ($h_s$): Encodes static or external environmental information.
   • Agent Features ($h_a$): Encodes the configuration of the agent.
   • Constraint: The dimension of $h_a$ is constrained to be a small fraction of the total latent space, creating a bottleneck that forces the model to focus only on the most predictable dynamics (the agent itself).
3. Prediction: A recurrent network predicts the future agent features ($\hat{h}_a$) based on the current agent features and the sequence of actions.
4. Reconstruction: The predicted future agent features are concatenated with the frozen scene features (from the context) and passed to a decoder to reconstruct the future frame ($\hat{x}_k$).

Training Objective

The model is trained to minimize the reconstruction error between the predicted future frame and the actual future frame:
$$L_{ef} = \mathbb{E}_{\tau \sim D} [\|\hat{x}_k - x_k\|_2^2]$$
Because the scene features are "fast-forwarded" from the context, the model is forced to learn the dynamics of the agent to reconstruct the moving parts of the image. If the agent picks up a tool, the tool's motion becomes predictable from the agent's actions, and the model learns to include the tool in the agent representation.

Integration with RL

EF is integrated as an auxiliary task into existing RL algorithms (specifically DrQ-v2 for model-free and TD-MPC2 for model-based).

• Shared Encoder: The visual encoder is shared between the RL policy/critic and the EF prediction task.
• Joint Optimization: The total loss is a weighted sum of the RL loss and the EF reconstruction loss:
  $$L_{total} = L_{RL} + \beta L_{ef}$$
• Motor Babbling: To ensure the agent explores diverse movements during the early stages (when it might otherwise just maximize reward in a limited way), a "motor babbling" phase is introduced where actions are random, forcing the learning of the visuomotor mapping.

3. Key Contributions

1. Self-Supervised Disentanglement: A novel method to separate agent and environment representations without requiring segmentation masks or external supervision.
2. Adaptability to Body Schema Changes: The method successfully adapts to dynamic changes in the agent's embodiment, such as picking up and using tools, by treating the tool as part of the predictable agent dynamics.
3. Improved Sample Efficiency: Demonstrated that adding EF as an auxiliary task significantly improves the sample efficiency and asymptotic performance of state-of-the-art RL algorithms.
4. General Applicability: The approach is architecture-agnostic and was successfully tested on both model-free (DrQ-v2) and model-based (TD-MPC2) algorithms.

4. Experimental Results

Experiments were conducted on two challenging benchmarks: Meta-World (robotic manipulation) and DMC (locomotion and manipulation), implemented in MuJoCo.

• Disentanglement Validation:
  • Visualization of gradients showed that as training progresses, the agent features ($h_a$) specialize in the robot's body and tools, while scene features ($h_s$) handle the background and static objects.
  • In "Hammer" tasks, the model correctly predicted the hammer's motion as part of the agent once grasped. In "Door Open" tasks, the door remained static in the prediction (treated as environment) because its motion was not fully determined by the agent's actions.
• RL Performance:
  • DrQ-v2-EF: Outperformed the baseline DrQ-v2 in 21 out of 26 tasks. It significantly reduced the number of steps required to reach 95% of maximum performance.
  • Tool Use: The performance gap between EF-augmented models and baselines was largest in tasks requiring tool use, validating the method's ability to integrate tools into the agent's body schema.
  • TD-MPC2-EF: Improved performance in 8 out of 10 tasks compared to the original TD-MPC2.
  • Comparison: The EF-augmented models outperformed the supervised baseline SEAR and the contrastive learning baseline CURL, and were competitive with the state-of-the-art Dreamer-v3.
• Ablation Studies:
  • The weight $\beta$ of the EF loss acts as a regularizer; too little has no effect, too much degrades performance.
  • The dimensionality of agent features must be small; larger dimensions allow the model to "cheat" by predicting environment dynamics, reducing disentanglement.
  • A prediction horizon ($H$) of 10 frames was found to be optimal for balancing performance and computational cost.

5. Significance and Conclusion

Ego-Foresight represents a significant step toward real-world robotic learning by removing the dependency on supervised segmentation masks, which are often unavailable or impractical in physical environments.

• Theoretical Impact: It validates the hypothesis that motor prediction is a powerful mechanism for learning agent-aware representations, mirroring human developmental processes.
• Practical Impact: By improving sample efficiency, EF reduces the computational and time costs of training RL agents. Its ability to adapt to new tools without retraining or re-supervision makes it highly suitable for general-purpose robotics.
• Future Directions: The authors suggest future work could replace pixel-wise reconstruction with contrastive losses to avoid modeling task-irrelevant details and explore scaling to high-resolution real-world data using Transformer architectures.

In summary, Ego-Foresight provides a robust, self-supervised framework that enables RL agents to "know themselves," leading to faster learning and better generalization in complex, dynamic environments.