PWM: Policy Learning with Multi-Task World Models

## Abstract

Abstract Reinforcement Learning (RL) has achieved impressive results on complex tasks but struggles in multi-task settings with different embodiments. World models offer scalability by learning a simulation of the environment, yet they often rely on inefficient gradient-free optimization methods. We introduce Policy learning with large World Models (PWM), a novel model-based RL algorithm that learns continuous control policies from large multi-task world models. By pre-training the world model on offline data and using it for first-order gradient policy learning, PWM effectively solves tasks with up to 152 action dimensions and outperforms methods using ground-truth dynamics. Additionally, PWM scales to an 80-task setting, achieving up to 27% higher rewards than existing baselines without the need for expensive online planning. Visualizations and code available at imgeorgiev.com/pwm

## 1 Introduction

The pursuit of generalizability in machine learning has recently been propelled by the training of large models on substantial datasets (brown2020language; kirillov2023segment; bommasani2021opportunities). Such advancements have notably permeated robotics, where multi-task behavior cloning techniques have shown remarkable performance (brohan2023rt; team2023octo; goyal2023rvt; bousmalis2023robocat). Nevertheless, these approaches predominantly hinge on near-expert data and struggle with adaptability across diverse robot morphologies due to their dependence on teleoperation (brohan2023rt; team2023octo; kumar2021should).

In contrast, Reinforcement Learning (RL) offers a robust framework capable of learning from suboptimal data, addressing the aforementioned limitations. However, RL has been focused on single-task experts (mnih2013playing; schulman2017proximal; haarnoja2018soft). Recent work suggests that a potential pathway to multi-task RL is with the world models framework, where a large model learns the environment dynamics and is then combined with zeroth-order planning to solve various tasks (ebert2018visual; schrittwieser2020mastering; xu2022xtra; feng2023finetuning; hansen2023td; schubert2023generalist). Notably, TD-MPC2 (hansen2023td) performs 80 continuous control tasks using a single 317M parameter multi-task world model.

Despite advancements, zeroth-order methods often result in suboptimal solutions due to high variance (mohamed2020monte; suh2022differentiable; parmas2023model) and online planning time scales with model size, rendering them infeasible at scale (schubert2023generalist; hansen2023td). These issues are particularly noticeable in world models with billions of parameters such as GAIA-1 (hu2023gaia) and UniSim (yang2023unisim). To address these challenges, we suggest a new approach: using world models as scalable, differentiable physics simulators to enable efficient policy learning through First-order Gradients (FoG). FoGs have demonstrated better performance and stability than traditional zeroth-order methods (xu2022accelerated). We analyze world models with differentiable simulation principles and find that the utility of world models lies not in their sheer accuracy but in (1) their smoothness, (2) their optimality gap, and (3) their gradient stability over long-horizon forecasts (suh2022differentiable; howell2022dojo).

Building on these insights, we introduce Policy learning with large World Models (PWM), a novel Model-Based RL (MBRL) algorithm and framework aimed at deriving effective continuous control policies from large, multi-task world models. We utilize pre-trained TD-MPC2 world models to efficiently learn control policies with FoG in <10m per task. Our empirical evaluations on complex locomotion tasks indicate that PWM not only achieves higher reward than baselines but also outperforms methods that use ground-truth simulation dynamics. In a multi-task scenario utilizing a pre-trained 48M parameter world model, PWM achieves up to 27% higher reward than TD-MPC2 without relying on online planning.

This underscores the efficacy of PWM and supports our broader contributions:

- 1.
Through pedagogical examples and ablations, we show that more accurate world models do not translate to better performing policies. Instead of pursuing world model improvements in isolation, we should aim to build world models that result in better policies.
- 2.
When regularized correctly, world models enable efficient first-order optimization. We show that this results in better performing policies and faster training times in comparison to zeroth-order gradient methods.
- 3.
We propose PWM, an algorithm for learning expert policies from large multi-task world models in minutes using First-order Gradients (FoG).

## 2 Related work

RL approaches can be classified as model-based and model-free which assume and do not assume a model respectively (arulkumaran2017deep). Most common algorithms for real-world applications such as PPO (schulman2017proximal) and SAC (haarnoja2018soft) are model-free and fall in the category of on-policy and off-policy methods respectively (arulkumaran2017deep). Both harness an actor-critic architecture where the critic learns the value and the actor optimizes the critic directly to maximize cumulative rewards (konda1999actor). Off-policy methods such as SAC typically learn a policy by taking First-order Gradients (FoG) directly from the critic. FoGs optimization typically has lower variance but is affected by discontinuities of the objective (mohamed2020monte). On-policy methods such as PPO harness zeroth-order gradients to learn policies (sutton1999policy). These gradients are not affected by discontinuities, making them effective for robotic tasks such as locomotion (rudin2022learning) but also exhibit high variance, which leads to slow optimization and suboptimal solutions (mohamed2020monte; suh2022differentiable). Differentiable simulation has been a popular framework to explore the properties of these gradient types (suh2022differentiable; howell2022dojo; xu2022accelerated).

While RL remains mostly focused on single-task policies, the larger robotics community has shifted focus to large multi-task models relying on behavior cloning (firoozi2023foundation). RT-1 (brohan2022rt) trains a large Transformer sequence model on a vast dataset to perform low-level object manipulation on a real robot. RT-2 (brohan2023rt) extends this by pre-training an even larger model on internet-scale data to also do planning and grasp affordance. Most recently, Open X (padalkar2023open) and Octo (team2023octo) showed that training these approaches on different tasks and embodiment leads to increased performance on single-task single-embodiment evaluations. Unfortunately, the benefits of large models and large datasets have not been sufficiently explored in the context of RL. DreamerV3 (hafner2023mastering) is a Model-Based Reinforcement Learning (MBRL) approach which learns a representation of its environment with encoder, decoder, dynamics and reward modules (collectively dubbed world model). It showed that scaling to larger models (400M parameters) leads to higher data efficient and episode rewards. TD-MPC2 (hansen2023td) continued this trend by employing a 317M parameter world model to learn an online planning multi-task policy capable of achieving good performance on 80 tasks. While showing impressive multi-task scalability, TD-MPC2 fails to solve all tasks it was trained on and exhibits limited scalability due to its use of online planning. With world models increasingly growing in size, we need a policy learning technique capable of scaling to billion parameter world models such as GAIA-1 (hu2023gaia) or UniSim (yang2023unisim).

## 3 Background

We focus on discrete-time and infinite horizon Reinforcement Learning (RL) scenarios characterized by system states ${\bm{s}}\in\mathbb{R}^{n}=\mathcal{S}$, actions ${\bm{a}}\in\mathbb{R}^{m}=\mathcal{A}$, dynamics function $f:\mathcal{S}\times\mathcal{A}\rightarrow\mathcal{S}$ and a reward function $r:\mathcal{S}\times\mathcal{A}\rightarrow\mathbb{R}$. Combined, these form a Markov Decision Problem (MDP) summarized by the tuple $(\mathcal{S},\mathcal{A},f,r,\gamma)$ where $\gamma$ is the discount factor. Actions at each timestep $t$ are sampled from a stochastic policy ${\bm{a}}_{t}\sim\pi_{\theta}(\cdot|{\bm{s}}_{t})$, parameterized by ${\bm{\theta}}$. The goal of the policy is to maximize the cumulative discounted rewards:

$$ $\displaystyle\max_{\bm{\theta}}J({\bm{\theta}}):=\max_{\bm{\theta}} \operatorname{\mathbb{E}}_{\begin{subarray}{c}{\bm{s}}_{1}\sim\rho(\cdot)\\ {\bm{a}}_{t}\sim\pi_{\bm{\theta}}(\cdot|{\bm{s}}_{t})\end{subarray}}\left[\sum _{t=1}^{\infty}\gamma^{t}r({\bm{s}}_{t},{\bm{a}}_{t})\right]$ (1) $$

where $\rho({\bm{s}}_{1})$ is the initial state distribution. Since this maximization over an infinite sum is intractable, in practice we often maximize over a value estimate. The value of a state ${\bm{s}}_{t}$ is defined as the expected reward follow the policy $\pi_{\bm{\theta}}$

$$ $\displaystyle V^{\pi}_{\bm{\psi}}({\bm{s}}_{t}):=\operatorname{\mathbb{E}}_{{ \bm{a}}_{h}\sim\pi_{\bm{\theta}}(\cdot|{\bm{s}}_{h})}\left[\sum_{h=t}^{\infty} \gamma^{h}r({\bm{s}}_{h},{\bm{a}}_{h})\right]$ (2) $$

When $V$ is approximated with a learned model with parameters ${\bm{\psi}}$ and $\pi_{\bm{\theta}}$ attempts to maximize some function of $V$, we arrive at the popular and successful actor-critic architecture (konda1999actor). Additionally, in MBRL it is common to also learn approximations of $f$ and $r$, which we denote as $F_{\bm{\phi}}$ and $R_{\bm{\phi}}$, respectively. It has also been shown to be beneficial to encode the true state ${\bm{s}}$ into a latent state ${\bm{z}}$ using a learned encoder $E_{\bm{\phi}}$ (hafner2019dream; hansen2022temporal; hansen2023td; hafner2023mastering). Putting together all of these components we can define a model-based actor-critic algorithm to consist of the tuple $(\pi_{\bm{\theta}},V_{\bm{\psi}},E_{\bm{\phi}},F_{\bm{\phi}},R_{\bm{\phi}}$) which can describe popular approaches such as Dreamer (hafner2019dream) and TD-MPC2 (hansen2023td). Notably, we make an important distinction between the types of components. We refer to $E_{\bm{\phi}}$, $F_{\bm{\phi}}$ and $R_{\bm{\phi}}$ as the world model components since they are a supervised learning problem with fixed targets. On the other hand, $\pi_{\theta}$ and $V_{\psi}$ optimize for moving targets which is fundamentally more challenging and we refer to them as the policy components.

## 4 Policy optimization through learned world models

This paper builds on the insight that since access to $F_{\bm{\phi}}$ and $R_{\bm{\phi}}$ is assumed through a pre-trained world-model, we have the option to optimize Eq. 1 via First-order Gradient (FoG) optimization which exhibit lower gradient variance, more optimal solutions and improved sample efficiency (mohamed2020monte). In our setting, these types of gradients are obtained by directly differentiating the expected terms of Eq. 1 as shown in Eq. 3. Note that this gradient estimator is also known as reparameterized gradient (kingma2015variational) and pathwise derivative (schulman2015gradient). While we use the explicit $\nabla^{[1]}$ notation below, we later drop it for simplicity as all gradient types in this work are first-order gradients.

$$ $\displaystyle\nabla_{{\bm{\theta}}}^{[1]}J({\bm{\theta}}):=\operatorname{ \mathbb{E}}_{\begin{subarray}{c}{\bm{s}}_{1}\sim\rho(\cdot)\\ {\bm{a}}_{h}\sim\pi_{\bm{\theta}}(\cdot|{\bm{s}}_{h})\end{subarray}}\left[ \nabla_{\bm{\theta}}\bigg{(}\sum_{t=1}^{\infty}\gamma^{t}r({\bm{s}}_{t},{\bm{a }}_{t})\bigg{)}\right]$ (3) $$

As $\nabla_{{\bm{\theta}}}^{[1]}J({\bm{\theta}})$ in itself is a random variable, we need to estimate it. A popular way to do that in practice is via Monte-Carlo approximation where we are interested in two properties - bias and variance. In Sections 4.1 and LABEL:sec:learning-chaos we tackle each aspect with toy robotic control problem to build intuition. In Section LABEL:sec:method we combine our findings to propose a new algorithm.

### 4.1 Learning through contact

Figure: (a) Ball-wall visualization.
Refer to caption: x3.png

## Peer Reviews

confidence: 4
contribution: 2
initial_rating: 5
presentation: 3
questions: Please see the weaknesses section.
soundness: 3
strengths: 1. This paper is clearly written and easy to follow. 
2. This paper presents sufficient experimental results to demonstrate the validity of its proposed method.
summary: This paper introduces a novel model-based reinforcement learning approach to pretrain a large multi-task world model on offline data with multiple tasks, followed by learning policies using first-order optimization for each task individually. Extensive experiments on DMC and Metaworld are conducted to showcase the strengths of the proposed method. While claiming to tackle multi-task challenges, this work tends to train a singular strategy for each isolated task.
weaknesses: 1. My primary concern revolves around the novelty of the proposed approach. This study appears to amalgamate TD-MPC2 and SHAC methodologies. To elaborate, a multi-task world model incorporating a task embedding and SimNorm activation mirrors aspects of TD-MPC2 for representation learning. Subsequently, during the policy learning phase, policies are acquired via first-order optimization akin to SHAC. Consequently, the original contribution of this work may seem somewhat limited.
2. In Figure 13, PWM exhibits inferior performance compared to SHAC in high-dimensional tasks like Anymal, Humanoid, and SHU Humanoid. Does this observation suggest that in intricate environments, precise world models hold greater significance?
3. DreamerV3 is a powerful baseline model capable of handling various domains with a single configuration. It would be best to compare the proposed method with it in Figure 5.
4. In Figure 6, the left figure indicates that TM-MPC2 outperforms PWM in Metaworld, contradicting the statement "In both settings PWM achieves higher reward than TD-MPC2 without the need for online planning."
5. In Lines 348-349, the mean of 'm' and 'n' are not explicitly defined or mentioned.

---

confidence: 3
contribution: 3
initial_rating: 6
presentation: 3
questions: 1) Figure 6, MT80: is something wrong with this plot on the left? It look like PWM performs WAY worse, or am I misreading it? I am assuming this was a drafting error, but if it is not then it deserves more discussion.

2) The use of V in Eqn 4 is confusing, because I don’t think it has any relation to Eqn 2 (or does it?)

3) What are scenarios (if any) where this reduction in world model accuracy for the sake of smoothness could be problematic? I would appreciate a little more discussion. Perhaps for transfer between tasks that share dynamics but do not share any reward structure? Does the discrepancy between TDMPC2 and PWM on mw-stick-pull potentially point to this? 

4) What might explain the reduced performance of PWM on the higher dimensional tasks, relative to SHAC (Figure 13)? Inaccuracies of the world model? If so, perhaps due to not having enough training data, or due to the regularization of the world model?

5) Where does the dataset come from? Is a source algorithm such as PPO used to gather the data? How much do these results depend on that? Do the datasets have examples of successful performance on each task? Does data from the task for which a policy is extracted need to have been included in the initial world model training?

7) Why can’t there be a comparison with DreamerV3 in a multi-task setting? In general, I am somewhat confused about how Dreamer, which I believe also uses first-order gradient optimization through a differentiable world model, relates to this work and would appreciate a little more discussion. 

8) Is performance sensitive to the simnorm hyper parameters?


Minor: 
- Some initial confusions (that were eventually cleared up, but hung around in my head for a while): what does '10 minutes' refer to? Of gameplay? Computation? On one GPU? Is new data acquired per task or it is reusing the pre-trained offline data for each specific task?
- I had to look at the specific references to understand that the first sentence of the intro was talking about foundation models. It is an interesting comparison, and I like the point, but could perhaps be made a little bit more obvious.
- In discussing large-data approaches to multi-task performance, I wonder how this relates to SIMA from the paper ‘Scaling Instructable Agents Across Many Simulated Worlds’?
- Typo in caption of Fig 1 (‘prpose’)
- Typo line 188 (‘optimality hap’)
- Typo line 377 (‘generalizble’)
soundness: 3
strengths: - Good baseline comparisons with impressive results for PWM. The comparison to a method that uses ground truth dynamics gradients (from the differentiable simulator) is valuable and compelling. 

- Good ablation experiments investigating sensitivity to the environment dynamics and the degree of world model regularization.
summary: This paper learns a differentiable multi-task world model from offline data, and then learns a policy using first-order optimization through the world model. They find that a well-regularized world model (they use SimNorm activation to normalize the latent vector) that prioritizes smoothness over correctness of the transition predictions is more effective for learning policies. They explain this phenomenon with two pedagogical examples, related to discontinuous contact dynamics and chaotic dynamics. They demonstrate improvement over TD-MPC2, a leading MBRL agent published at the previous ICLR, and in some cases over SHAC, an approach that relies on ground truth world model differentiation.
weaknesses: - A diagram of the model and policy training would be helpful (i.e. more detailed than Figure 1). Perhaps specifically illustrating what is different about this vs. TD-MPC2. 

- I find Figure 2 confusing: What is the y-axis of the middle plot? What does a negative value mean? Is theta the angle (perhaps confusing because it is overloading the use of the theta symbol)? Additionally, I’m a little confused about how this example relates to a broader phenomenon of contact-induced discontinuities: the main problem here appears to be the long stretch of flat J, which causes the optimization to get stuck, as opposed to the fact that there is a non-differentiable discontinuity. Is it just a fluke of this toy scenario that the regularization helps? Why does SimNorm decrease the optimality gap? I’m a little ambivalent as to whether this example helps with understanding or just adds confusion, and on my initial read-throughs I didn’t find this “intuitive” example to be super intuitive, at least as-presented.

---

confidence: 4
contribution: 3
initial_rating: 5
presentation: 2
questions: 1. Since there are no interactions between the target policy and the environment, could this method be considered as an offline policy learning method? If so, how would this method solve the distributional shift problem? For example, will the world overestimate the value of some state-action pairs not included in the dataset and result in poor performance?
2. As mentioned in weakness 2, what will you consider to be the major differences between your method and MAAC? 
3. I would like to know the comparison of the total training time (including world model training) of the PWM with other model-free methods like PPO and SAC.
4. How is "high-dimensional" defined in your work? For example, in Figure 4, "Hopper" is also considered a high-dimensional task. But according to the document (https://gymnasium.farama.org/environments/mujoco/hopper/), it only has an 11-dim observation space and a 3-dim action space.
5. According to Table 3 and Table 4, the PWM method shows much larger variances than model-free methods. What's the main reason for this and is it able to avoid the problem?
6. Some curves in Figure 13 are incomplete. What's the reason for that?

I am happy to increase my score if the authors justify these questions.
soundness: 3
strengths: 1. The key idea of the paper is well delivered using pedagogical examples and is easy to follow

2. The results are highly reproducible with detailed experimental settings and code included, which also benefits the community for further study.

3. The experiments conducted in this paper are quite comprehensive and solid.
summary: This paper introduces Policy learning with multi-task World Models (PWM) as a new model-based reinforcement learning algorithm. The main idea is to employ a pre-trained regularized and differentiable world model to simulate the environmental dynamics and then adopt first-order gradient methods to obtain the optimal policy efficiently. The author conducted experiments on both contact-rich single-task and large multi-task settings, showing that PWM had a better performance in learning high-reward policies and drastically improved the training time compared to previous methods. The author also provides additional examples and ablation studies to show the correlation between the accuracy of world models and the final performance, revealing that smoothness instead of accuracy is the key factor for higher rewards in gradient optimization.
weaknesses: 1. Typos and incorrect format
- Incorrect citing format in chapter 4.2: when the authors or the publication are included in the sentence, the citation should not be in parenthesis, please use \citet{} instead.
- Typos in Equation (2): it should be $\gamma^{h-t+1}$ instead of $\gamma^h$
- Typos in Algorithm 1: "$\nabla$" is omitted in all equations regarding learnable parameters update.
- Typos in Figure 10: missing ")" in the caption.

2. Lack of novelty
- The key idea of policy optimization via the pathwise derivative of a learned world model has already been proposed in MAAC[1], which is not included in the related works. The author should show the main differences between MAAC and the proposed PWM via analysis or experiments.

3. A total training time cost comparison between PWM and other algorithms is not revealed in the paper.

[1] Clavera, Ignasi, Yao Fu, and Pieter Abbeel. "Model-Augmented Actor-Critic: Backpropagating through Paths." International Conference on Learning Representations (2020).

---

confidence: 5
contribution: 2
initial_rating: 3
presentation: 2
questions: Besides, there are several minor questions.

1. For the single-task setting, what's the difference between PWM and Dreamer (beyond implementation details)? Can PWM compare to DreamerV3 in the setting of Section 4.1?
2. In Line 298, do you mean world models are also finetuned with online interaction data? How about the multi-task setting? It seems the multi-task setting is fully offline.
3. In Figure 6 (left), why PWM significantly underperforms TD-MPC2?
4. In Line 357, does the baseline TD-MPC2 use the same hyperparameter as $H=16$?
5. Can the author explain why PWM benefits from smaller batch sizes?
6. Typo: Line 036 prpose => propose

Overall, I find this paper interesting and informative. If the authors address my concerns well, I am willing to raise my rating.
soundness: 3
strengths: 1. Didactic examples are informative.
2. Experiment results: sample efficiency to learn a task within 10 minutes seems promising.
3. The paper is well-written and easy to understand.
summary: This paper introduces PWM, proposing to use first-order gradient optimization for policy learning on top of offline pretrained world models. This method differs from zero-order optimization (MPPI) used in its base method TD-MPC2. Two didactic examples are presented to demonstrate that a well-regularized world model may be more useful for policy learning compared to real dynamics, confirming the well-known objective mismatch problem in the community. Main experimetns are conducted for single-task and multi-task learning, outperforming model-free baselines, zero-order method TD-MPC2, and method using real dynamics.
weaknesses: The major concerns are novelty and clarity of contributions. This paper does not well highlight its own contribution to the existing literature.

- The paper's contribution is conveyed (at least partially) on top of the fact that "world models methods often rely on inefficient gradient-free optimization methods for policy extraction." This holds for the TD-MPC series of work but is not valid for another series of popular methods, Dreamer. Actually, since DreamerV1, the policy learning process is done by first-order gradient optimization using the world model as a differentiable dynamics. The formulation of actor-critic learning in Equ (6-9) is almost the same as in the Dreamer series.
- Based on the above discussion, I believe the contribution of this work is additions to the following two problems: 1. objective mismatch in MBRL; 2. a multi-task RL with world models. However, contributions to this two aspect are also not groundbreaking. For the former, this problem has been well-known since Lambert et al. (2020), and for the latter, MBRL is also well-known for being beneficial for multi-tasking.

Meta-review: This paper argues for the utility of well-regularized world-models for first-order optimization and propose Policy learning with multi-task World Models (PWM), consisting of pre-training a world-model on offline data and then extracting policies from it using first-order optimization. The authors present an empirical evaluation over a large number of tasks, comparing against related baselines.

The paper is well-written, explains concepts quite well (via some pedagogical examples), and provides strong empirical evaluations.

The main weakness of this work is its novelty. While the method is novel and results in strong empirical performance, none of its component parts (training world models, using first-order gradient optimization, and training on multi-tasks) is novel. The analysis in sections 3.1 and 3.2 is quite nice, but feel somewhat limited for one of the main claims of the paper ("Correlation Between World Model Smoothness and Policy Performance").

While their analysis does motivate their use in FoG optimization, the use of multiple tasks for training seems to come unnanounced. The fact that PWM works well in the multi-task setting does not answer _why_ the authors felt it necessary to add this. For instance, the pedagogical examples were all single-task settings, and this analysis was used to motivate the use of world-models for PWM.

In summary, this is an interesting idea that is well-presented, but I don't believe it is yet at its full potential. In particular, the correlation between world-model smoothness and policy performance bears digging into, as these insights will likely prove more useful to the general community than PWM itself. As such, I would personally recommend the authors extend this work with more analyses of that type so as to strengthen each of the three main claims.

However, given that the reviewers are unanimous in being supportive of acceptance, I will not override their decision.

Decision: Accept (Poster)