Improving World Models using Deep Supervision with Linear Probes

• URL: http://arxiv.org/abs/2504.03861v1
• Date: Fri, 04 Apr 2025 18:35:21 GMT
• Title: Improving World Models using Deep Supervision with Linear Probes
• Authors: Andrii Zahorodnii,
• Abstract summary: In this paper, we investigate a deep supervision technique for encouraging the development of a world model in a network trained end-to-end to predict the next observation.<n>Using an experimental environment based on the Flappy Bird game, we explore the effect of adding a linear probe component to the network's loss function.<n>Our experiments demonstrate that this supervision technique improves both training and test performance, enhances training stability, and results in more easily decodable world features.
• Score: 0.0
• License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/
• Abstract: Developing effective world models is crucial for creating artificial agents that can reason about and navigate complex environments. In this paper, we investigate a deep supervision technique for encouraging the development of a world model in a network trained end-to-end to predict the next observation. While deep supervision has been widely applied for task-specific learning, our focus is on improving the world models. Using an experimental environment based on the Flappy Bird game, where the agent receives only LIDAR measurements as observations, we explore the effect of adding a linear probe component to the network's loss function. This additional term encourages the network to encode a subset of the true underlying world features into its hidden state. Our experiments demonstrate that this supervision technique improves both training and test performance, enhances training stability, and results in more easily decodable world features -- even for those world features which were not included in the training. Furthermore, we observe a reduced distribution drift in networks trained with the linear probe, particularly during high-variability phases of the game (flying between successive pipe encounters). Including the world features loss component roughly corresponded to doubling the model size, suggesting that the linear probe technique is particularly beneficial in compute-limited settings or when aiming to achieve the best performance with smaller models. These findings contribute to our understanding of how to develop more robust and sophisticated world models in artificial agents, paving the way for further advancements in this field.

Related papers

• AI in a vat: Fundamental limits of efficient world modelling for agent sandboxing and interpretability [84.52205243353761]
  Recent work proposes using world models to generate controlled virtual environments in which AI agents can be tested before deployment. We investigate ways of simplifying world models that remain agnostic to the AI agent under evaluation.
  arXiv  Detail & Related papers  (2025-04-06T20:35:44Z)
• AdaWorld: Learning Adaptable World Models with Latent Actions [76.50869178593733]
  We propose AdaWorld, an innovative world model learning approach that enables efficient adaptation.<n>Key idea is to incorporate action information during the pretraining of world models.<n>We then develop an autoregressive world model that conditions on these latent actions.
  arXiv  Detail & Related papers  (2025-03-24T17:58:15Z)
• Robotic World Model: A Neural Network Simulator for Robust Policy Optimization in Robotics [50.191655141020505]
  We introduce a novel framework for learning world models.<n>By providing a scalable and robust framework, we pave the way for adaptive and efficient robotic systems in real-world applications.
  arXiv  Detail & Related papers  (2025-01-17T10:39:09Z)
• Learning World Models for Unconstrained Goal Navigation [4.549550797148707]
  We introduce a goal-directed exploration algorithm, MUN, for learning world models. MUN is capable of modeling state transitions between arbitrary subgoal states in the replay buffer. Results demonstrate that MUN strengthens the reliability of world models and significantly improves the policy's capacity to generalize.
  arXiv  Detail & Related papers  (2024-11-03T01:35:06Z)
• Open-World Reinforcement Learning over Long Short-Term Imagination [91.28998327423295]
  Training visual reinforcement learning agents in a high-dimensional open world presents significant challenges.<n>We present LS-Imagine, which extends the imagination horizon within a limited number of state transition steps.<n>Our method demonstrates significant improvements over state-of-the-art techniques in MineDojo.
  arXiv  Detail & Related papers  (2024-10-04T17:17:30Z)
• ReCoRe: Regularized Contrastive Representation Learning of World Model [21.29132219042405]
  We present a world model that learns invariant features using contrastive unsupervised learning and an intervention-invariant regularizer. Our method outperforms current state-of-the-art model-based and model-free RL methods and significantly improves on out-of-distribution point navigation tasks evaluated on the iGibson benchmark.
  arXiv  Detail & Related papers  (2023-12-14T15:53:07Z)
• Predictive Experience Replay for Continual Visual Control and Forecasting [62.06183102362871]
  We present a new continual learning approach for visual dynamics modeling and explore its efficacy in visual control and forecasting. We first propose the mixture world model that learns task-specific dynamics priors with a mixture of Gaussians, and then introduce a new training strategy to overcome catastrophic forgetting. Our model remarkably outperforms the naive combinations of existing continual learning and visual RL algorithms on DeepMind Control and Meta-World benchmarks with continual visual control tasks.
  arXiv  Detail & Related papers  (2023-03-12T05:08:03Z)
• DITTO: Offline Imitation Learning with World Models [21.419536711242962]
  DITTO is an offline imitation learning algorithm which addresses all three of these problems.<n>We optimize this multi-step latent divergence using standard reinforcement learning algorithms.<n>Our results show how creative use of world models can lead to a simple, robust, and highly-performant policy-learning framework.
  arXiv  Detail & Related papers  (2023-02-06T19:41:18Z)
• Bridging the Gap to Real-World Object-Centric Learning [66.55867830853803]
  We show that reconstructing features from models trained in a self-supervised manner is a sufficient training signal for object-centric representations to arise in a fully unsupervised way. Our approach, DINOSAUR, significantly out-performs existing object-centric learning models on simulated data.
  arXiv  Detail & Related papers  (2022-09-29T15:24:47Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.