PWM: Policy Learning with Multi-Task World Models

• URL: http://arxiv.org/abs/2407.02466v3
• Date: Mon, 24 Feb 2025 06:56:00 GMT
• Title: PWM: Policy Learning with Multi-Task World Models
• Authors: Ignat Georgiev, Varun Giridhar, Nicklas Hansen, Animesh Garg,
• Abstract summary: World model methods offer scalability by learning a simulation of the environment.<n> gradient-based methods exhibit lower variance but fail to handle discontinuities.<n>We introduce Policy learning with multi-task World Models (PWM), a novel model-based RL algorithm for continuous control.
• Score: 37.678858748473196
• License: http://creativecommons.org/licenses/by/4.0/
• Abstract: Reinforcement Learning (RL) has made significant strides in complex tasks but struggles in multi-task settings with different embodiments. World model methods offer scalability by learning a simulation of the environment but often rely on inefficient gradient-free optimization methods for policy extraction. In contrast, gradient-based methods exhibit lower variance but fail to handle discontinuities. Our work reveals that well-regularized world models can generate smoother optimization landscapes than the actual dynamics, facilitating more effective first-order optimization. We introduce Policy learning with multi-task World Models (PWM), a novel model-based RL algorithm for continuous control. Initially, the world model is pre-trained on offline data, and then policies are extracted from it using first-order optimization in less than 10 minutes per task. PWM effectively solves tasks with up to 152 action dimensions and outperforms methods that use ground-truth dynamics. Additionally, PWM scales to an 80-task setting, achieving up to 27% higher rewards than existing baselines without relying on costly online planning. Visualizations and code are available at https://www.imgeorgiev.com/pwm/.

Related papers

• Policy-Driven World Model Adaptation for Robust Offline Model-based Reinforcement Learning [6.189693079685375]
  offline model-based RL (MBRL) explicitly learns a world model from a static dataset.<n>We propose a framework that dynamically adapts the world model alongside the policy.<n>We benchmark our algorithm on twelve noisy D4RL MuJoCo tasks and three Tokamak Control tasks, demonstrating its state-of-the-art performance.
  arXiv  Detail & Related papers  (2025-05-19T20:14:33Z)
• Learning from Reward-Free Offline Data: A Case for Planning with Latent Dynamics Models [79.2162092822111]
  We systematically evaluate reinforcement learning (RL) and control-based methods on a suite of navigation tasks.<n>We employ a latent dynamics model using the Joint Embedding Predictive Architecture (JEPA) and employ it for planning.<n>Our results show that model-free RL benefits most from large amounts of high-quality data, whereas model-based planning generalizes better to unseen layouts.
  arXiv  Detail & Related papers  (2025-02-20T18:39:41Z)
• MOTO: Offline Pre-training to Online Fine-tuning for Model-based Robot Learning [52.101643259906915]
  We study the problem of offline pre-training and online fine-tuning for reinforcement learning from high-dimensional observations. Existing model-based offline RL methods are not suitable for offline-to-online fine-tuning in high-dimensional domains. We propose an on-policy model-based method that can efficiently reuse prior data through model-based value expansion and policy regularization.
  arXiv  Detail & Related papers  (2024-01-06T21:04:31Z)
• Gradient-based Planning with World Models [21.9392160209565]
  We present an exploration of a gradient-based alternative that fully leverages the differentiability of the world model. In a sample-efficient setting, our method achieves on par or superior performance compared to the alternative approaches in most tasks.
  arXiv  Detail & Related papers  (2023-12-28T18:54:21Z)
• Deep Model Predictive Optimization [21.22047409735362]
  A major challenge in robotics is to design robust policies which enable complex and agile behaviors in the real world. We propose Deep Model Predictive Optimization (DMPO), which learns the inner-loop of an MPC optimization algorithm directly via experience. DMPO can outperform the best MPC algorithm by up to 27% with fewer samples and an end-to-end policy trained with MFRL by 19%.
  arXiv  Detail & Related papers  (2023-10-06T21:11:52Z)
• Reparameterized Policy Learning for Multimodal Trajectory Optimization [61.13228961771765]
  We investigate the challenge of parametrizing policies for reinforcement learning in high-dimensional continuous action spaces. We propose a principled framework that models the continuous RL policy as a generative model of optimal trajectories. We present a practical model-based RL method, which leverages the multimodal policy parameterization and learned world model.
  arXiv  Detail & Related papers  (2023-07-20T09:05:46Z)
• Model-Based Reinforcement Learning with Multi-Task Offline Pretraining [59.82457030180094]
  We present a model-based RL method that learns to transfer potentially useful dynamics and action demonstrations from offline data to a novel task. The main idea is to use the world models not only as simulators for behavior learning but also as tools to measure the task relevance. We demonstrate the advantages of our approach compared with the state-of-the-art methods in Meta-World and DeepMind Control Suite.
  arXiv  Detail & Related papers  (2023-06-06T02:24:41Z)
• Exploration via Planning for Information about the Optimal Trajectory [67.33886176127578]
  We develop a method that allows us to plan for exploration while taking the task and the current knowledge into account. We demonstrate that our method learns strong policies with 2x fewer samples than strong exploration baselines.
  arXiv  Detail & Related papers  (2022-10-06T20:28:55Z)
• Visual Foresight With a Local Dynamics Model [1.370633147306388]
  We propose the Local Dynamics Model (LDM) which efficiently learns the state-transition function for single-step manipulation primitives. By combining the LDM with model-free policy learning, we can learn policies which can solve complex manipulation tasks using one-step lookahead planning.
  arXiv  Detail & Related papers  (2022-06-29T17:58:14Z)
• Temporal Difference Learning for Model Predictive Control [29.217382374051347]
  Data-driven model predictive control has two key advantages over model-free methods. TD-MPC achieves superior sample efficiency and performance over prior work on both state and image-based continuous control tasks.
  arXiv  Detail & Related papers  (2022-03-09T18:58:28Z)
• Evaluating model-based planning and planner amortization for continuous control [79.49319308600228]
  We take a hybrid approach, combining model predictive control (MPC) with a learned model and model-free policy learning. We find that well-tuned model-free agents are strong baselines even for high DoF control problems. We show that it is possible to distil a model-based planner into a policy that amortizes the planning without any loss of performance.
  arXiv  Detail & Related papers  (2021-10-07T12:00:40Z)
• Evolutionary Planning in Latent Space [7.863826008567604]
  Planning is a powerful approach to reinforcement learning with several desirable properties. We learn a world model that enables Evolutionary Planning in Latent Space. We show how to build a model of the world by bootstrapping it with rollouts from a random policy and iteratively refining it with rollouts from an increasingly accurate planning policy.
  arXiv  Detail & Related papers  (2020-11-23T09:21:30Z)
• TempLe: Learning Template of Transitions for Sample Efficient Multi-task RL [18.242904106537654]
  TempLe is the first PAC-MDP method for multi-task reinforcement learning. We present two algorithms for an "online" and a "finite-model" setting respectively. We prove that our proposed TempLe algorithms achieve much lower sample complexity than single-task learners or state-of-the-art multi-task methods.
  arXiv  Detail & Related papers  (2020-02-16T19:46:49Z)
• Information Theoretic Model Predictive Q-Learning [64.74041985237105]
  We present a novel theoretical connection between information theoretic MPC and entropy regularized RL. We develop a Q-learning algorithm that can leverage biased models.
  arXiv  Detail & Related papers  (2019-12-31T00:29:22Z)

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.