PWM: Policy Learning with Multi-Task World Models
Authors: Ignat Georgiev
Varun Giridhar
Nicklas Hansen
Animesh Garg
Venue: ICLR 2025 Poster (2025)
Keywords: reinforcement learning
model-based reinforcement learning
continuous control
world models
Primary area: reinforcement learning
TL;DR: We propose an MBRL method that leverages pre-trained, multi-task world models for efficient policy learning using first-order optimization.

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 [imgeorgiev.com/pwm](https://www.imgeorgiev.com/pwm/).

Published as a conference paper at ICLR 2025
PWM: POLICY LEARNING WITH MULTI-TASK WORLD MODELS
Ignat Georgiev Varun Giridhar's ' Nicklas Hansen 2, Animesh Garg'
Georgia Institute of Technology
UC San Diego
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 generate smoother opti- mization 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 imgeorgiev. com/pwn
High-dim. single tasks
80 Multi-task
PPO-norm. Reward
Normalized Score
1.2
Policy
0.9
Large Multi-task
Policy
1.0
World
Policy
0.8
0.8
PPO
SHAC
DreamerV3
SAC
TD-MPC2
PWM (ours)
We propose PWM, method for multi-task RI. that utilizes pre-trained world models learn policies for cach task. When sufficiently regularized, these world models induce smooth optimization landscapes, which allows for efficient first-order optimization. Our approach can solve tasks in <10 minutes and achieves higher rewards both single-task and multi-task environments.
Figure 1:
INTRODUCTION
The pursuit of generalizability in machine learning has recently been propelled by the training of large models on substantial datasets Brown et al. (2020); Kirillov et al. (2023); Bommasani et al. (2021). Such advancements have notably permeated robotics, where multi-task behavior cloning techniques have shown remarkable performance Zitkovich et al. (2023); Octo Model Team et al. (2024); Goyal al. (2023); Bousmalis al. (2023). Nevertheless, these approaches predominantly hinge near-expert data and struggle with adaptability across diverse robot morphologies due to their dependence on teleoperation Zitkovich et al. (2023); Octo Model Team et al. (2024); Kumar al. (2021).
In contrast, Reinforcement Learning (RL) offers a robust framework capable learning from suboptimal data, addressing the aforementioned limitations. However, traditional RL has been focused on single-task experts Mnih et al. (2013); Schulman et al. (2017); Haarnoja et al. (2018). Recently, Hansen et al. (2024) suggested that potential pathway to multi-task RL is with the world

---

Published as a conference paper at ICLR 2025
models framework, where a large model learns the environment dynamics and is then combined with Zeroth-order Gradient (ZoG) methods. Despite advancements, ZoG methods struggle with sample inefficiency due to the high variance Mohamed et al. (2020); Suh et al. (2022); Parmas et al. (2023)
and online planning time scales with model size, rendering it infeasible at scale.
First-order Gradient (FoG) methods provide low-variance alternative that have shown superior sample efficiency and asymptotic performance when combined with smooth differentiable simulations Xu al. (2022). However, they struggle optimize through discontinuities Suh al. (2022); Georgiev et al. (2024). In this work, we explore the tight coupling between FoG optimization and world models through the lens of differentiable simulation. founter-intuitively we find that for gradient- based optimization, we don't 't want world models to be accurate; instead, we want them to be smooth and have a low optimality gap. This in turn enables efficient FoG optimization.
Building on these insights, we propose Policy learning with multi-task World Models (PWM), an algorithm that can learn policies from offline pre-trained world models in under <10 minutes per task. With this new-found efficiency, we also propose new multi-task framework, where instead of training a multi-task algorithm, we only train a multi-task world model and then extract a policy for each task. This decoupling of the supervised and RL objectives results in more stable and efficient learning with higher episode rewards. Our empirical evaluations on high-dim. tasks indicate that PWM not only achieves higher reward than baselines but also outperforms methods that use ground-truth dynamics. In a multi-task scenario utilizing a pre-trained 48M parameter world model from TD-MPC2, 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. Correlation Between World Model Smoothness and Policy Performance: Through
pedagogical examples and ablations, we demonstrate that smoother, better-regularized world models significantly enhance policy performance. Notably, this results in an inverse correlation between model accuracy and policy performance.
2. Efficiency of First-Order Gradient (FoG) Optimization: We show that combining FoG optimization with well-regularized world models enables more efficient policy learning compared zeroth-order methods. Furthermore, policies learned from world models asymp- totically outperform those trained with ground-truth simulation dynamics, emphasizing the importance of the tight relationship between FoG optimization and world model design.
3. Scalable Multi-Task Algorithm: Instead of training a single multi-task policy model, we propose PWM, a framework where a multi-task world model is first pre-trained on offline data. Then per-task expert policies are extracted in <10 minutes per task, offering clear and scalable alternative to existing methods focused on unified multi-task models.
2 BACKGROUND
We focus on discrete-time and infinite-horizon Reinforcement Learning (RL) scenarios characterized by system states E R" S, actions a E R A. dynamics function f S A - S and reward function : S x A + R. Combined, these form Markov Decision Problem (MDP) summarized by the tuple (S, A. f,r,) where 7 is the discount factor. Actions at each timestep are sampled from [o(- parameterized by O. The goal of the policy is to maximize the
a stochastic policy cumulative discounted rewards:
max J(Ø) == max E simp(-)
Σ En'r(s..a)) a.) 1
a.
where p(s)) is the initial state distribution. Since this maximization over an infinite sum is intractable, in practice we often maximize over value estimate. The value of state 8, is defined the expected reward following the policy
V (s:) == Taa)mm
Y (r(s, an)
Aett
When V is approximated with a learned model with parameters &b and To attempts to maximize some function of V. we arrive at the popular and successful actor-critic architecture Konda & Tsitsiklis (1999). Additionally, in MBRL it is common to also learn approximations of f and r, which we denote Fg and R2., respectively. It has also been shown to be beneficial encode the true state
2

---

Published as a conference paper at ICLR 2025
[(0)
20
ReLU
SimNorm
0
Model
Model entor
Opt.
gap
-20
True
0.0
16.85
y
ReLU
0.71
16.05
-2
0
2
SimNorm
1.13
3.47
(a) Ball-wall wall visualization.
(b) Problem landscape.
(c) Model error and optimality gap.
Figure 2: Ball-wall pedagogical example. The left figure visualizes the problem. The middle figure shows the problem landscape induced by cach model. J(0) shows the true underlying function and the two other are MLPs with different activation functions. We minimize each of these problems using gradient descent and starting at 0 = - (marker x). The colored crosses represent the solutions converged to for cach model. The right table shows the model approximation error during training and the optimality gap (J(0") - J(O) between the global minimum H and the solution found each model 0.
into latent state z using learned encoder Es Hafner al. (2019); Hansen et al. (2022; 2024); Hafner et al. (2023). Putting together all of these components, we can define a model- based actor- critic algorithm to consist of the tuple (To. Vas Eg. Fş. R&)) which can describe popular approaches such as Dreamer Hafner et al. (2019; 2023) and TD-MPC2 Hansen et al. (2024). Notably, we make an important distinction between the types of components. We refer to Eg,, Fe and Rgg as the world model components since they are supervised learning problems with fixed targets. On the other hand, Ty and Vo optimize for moving targets, which fundamentally more challenging, and we refer to them as the policy components.
POLICY OPTIMIZATION THROUGH WORLD MODELS
This paper builds on the insight that since access to Fs and Rs is assumed through a pre-trained world model, have the option optimize Eq. 1 via First-order Gradient (FoG) optimization which exhibit lower gradient variance, more optimal solutions, and improved sample efficiency Mohamed al. (2020). In setting, these types of gradients expected terms of Eq. as shown in Eq. 3. Note that this gradient estimator is also known as reparameterized gradient Kingma al. (2015) and pathwise derivative Schulman al. (2015). While the explicit 4 notation below, later drop for simplicity all gradient types in this work are first-order gradients.
obtained by directly differentiating the
((0)
= E #1 pimp(-) an amoTTo (-(()))))))
Vo Σ
rrr(s ai)
As (()(()) itself a random variable, we need to estimate it. A popular way to do that in practice via Monte-Carie approximation, where Sections 3.1 and 3.2 we tackle each aspect with a toy robotic control problem to build intuition. In algorithm.
interested in two properties: bias and variance. In
Section 3.3, combine findings propose
3.1 LEARNING THROUGH CONTACT
FoGs are unbiased E ( = E EE -() VJ((). only if both the dynamics f and rewards r are Lipschizz-smooth Suh et al. (2022). However, many robotic problems involving contact inherently non-smooth, which breaks these conditions and results in gradient sample error where E E(vÍ)]](0) * VJ((O) under finite number of samples N. Instead of directly optimizing the true, discontinuous objective, it is advantageous to optimize a smooth surrogate, such as a model learned by regularized deep neural network.
To illustrate this concept, we use a toy problem where a ball is thrown toward wall at a fixed velocity, shown in Figure 2a. The objective to find the optimal initial angle 0 such that maximize
3

---

Published as a conference paper at ICLR 2025
10-2
104
Var[v/(O)]
ESNR(V)(())
10-7
True
10
True
9
MLP H=3
MLP H=3
10
10-2
MLP H=16
MLP H=16
0 20
40
0
20
40
H
Figure 3: Double pendulum pedagogical example. The middle figure evaluates the variance of policy gradient estimates over N 100 Morle samples for varying horizons H. The right figure shows the same data but plots the Expected Signal-to-Nois ratio (ESNR) with higher values translating to more useful gradients. These results suggests that world models trained over long horizon trajectories provide more useful gradients. Note that H = 3 and H = 16 in the figure legends refer to the training horizon of the models.
forward distance. In this simplified pedagogical example, we assume that the ball "sticks" to the wall, creating discontinuous optimization landscape (Figure 2b). We compare the performance of two models in approximating this objective: a 2-layer Multi-Layer Perceptron (MLP) with ReLU activation and another MLP with SimNorm activation Hansen et al. (2024) in the intermediate layers. SimNorm normalizes latent vector z by projecting into simplices with dimension P using a softmax operator. Given an input vector z, SimNorm can be expressed mapping into L vectors:
SimNorm(2) (z) = [91..... gt].
Doftmax(2::444
(4)
We trained both MLPs and observed their effects smoothing the optimization landscape (Figure 2h). The ReLU-activated MLP smooths the landscape but introduces local minimum that hinders gradient descent, particularly when starting from between the solution and the optimal solution: ||l (l). In contrast, the SimNorm-activatte MLP has additional regularization, which reduces the optimality gap at the expense of model accuracy (Table 2c). This example highlights that more accurate models do not always lead to better policies, noted by Lambert et al. (2020). Our findings extend this by showing that for FoG optimization, prioritizing smoothness over accuracy can lead to improved results. Further details are provided in Appendix A.
-, resulting in large optimality gap (difference
3.2 LEARNING WITH CHAOTIC DYNAMICS
While FoGs have lower variance per step, they can accumulate significant variance over long-horizon rollouts Metz et al. (2021). Suh et al. (2022) link this variance to the smoothness of models and the length of the prediction horizon: Var [V] & [|V high variance renders FoGs ineffective in chaotic systems. Chaotic systems are characterized by their sensitivity to initial conditions, where small perturbations can lead to exponentially divergent trajectories, making long-term prediction particularly challenging. The double pendulum, also known as the Acrobot Murray & Hauser (1991), 1), is a classic example of such a system (Figure 3).
At sufficiently high H, the
We analyze the variance of gradient estimators in the double pendulum using both the true dynamics and a Simnorr--cttivated MLP model. The MLP model was trained for auto-regressive prediction horizons of И and H 16 until convergence on large dataset. Figure shows that both learned models exhibit reduced variance compared to the true dynamics. However, as noted by Parmas et al. (2023), variance alone is insufficient for drawing definitive conclusions about gradient quality. Instead, they propose analyzing gradients via their Expected Signal-to-Noise Ratio (ESNR), defined as:
ESNR(VJ(()) E VE(())(()) 2
(5)
Σ Var [
In Figure 3, we observe that learned models exhibit higher ESSR than the true dynamics, providing more useful gradients. Notably, the training horizon plays critical role, with the H = 16 model sustaining a higher ESNR over higher H. We conclude that learned world models offer more informative policy gradients than the true system dynamics. Further details in Appendix B.
4

---

Published as a conference paper at ICLR 2025
3.3 PWM: POLICY LEARNING WITH MULTI-TASK WORLD MODELS
Given the results from the previous subsection, we propose to view world models not as com- ponents of RL methods but instead as scalable, differentiable physies simulators that provide gradients with low sample error and variance. It is worth noting that approaches such as TD- MPC2 Hansen et (2024) do not exploit these properties but rather choose optimize policies via DtPG gradients: (oQ(s,
Algorithm 1: PWM: Policy learning with multi-task World Models
Given: Multi-task dataset Given: y: discount rate
Given: ag,
agg ag. learning rates
Initialize learnable parameters 0, 1b, Φ
b Pre-train world model once for N epochs do
VoJ(0) (-s)( (ooonti)i/Vo((ss.... a)].
81:H1 01:H TI-H, e Φ - Φ aΦ VL (Ø)
b Eq. 10
We propose a new method and framework for ef- ficiently learning policies from large multi-task world models.
end
Train policy on task embedding ng e for epochs do
Framework. Assuming availability of data from multiple tasks, we first train a multi-task world model to predict future states and rewards. Then for each task we want to solve, we learn single policy in minutes using FoG optimization. The policy is then deployed to solve the task and optionally fine-tune its world model and policy. Method. For policy learning, we propose on- policy actor-critic approach inspired by dif- ferentiable simulation approaches Xu et al. (2022) where the actor is trained via FoG back- propagated through the world model, while the critic is trained via TD(A). The key to our approach is that training is done in a batched fashion where multiple trajectories are imagined in parallel. The actor loss function is akin Eq. but features rewards over fixed horizon H, terminal value bootstrapping and usage of the learned world model components:
81 ~ B ž1 = E E)
for h=11, H] do
b Rollout
a₁ To (-12))
Th Rs Zh,, ah e) 1h+1 (z ah. e)
end
-- + ag VL, (0)
B Eq.
</ 1/ a+ VLv (36)
Eq. 7-9
end
H-1 Σ -h-1
C (0) = x1mp(-1
7 Ra an) + V (xH)
where
=1= (g(a
(6)
The critic is trained model-free fashion using TD(A) over an //-step rollout in latent space z as seen in other similar on-policy methods Sutton & Barto (2018); Hafner et al. (2019); Xu et al. (2022):
441-1
Vn(z,) := Σ
nn Ro(Znn and + V
n=t
H-1-1
(=z)== = (1- (1-
Σ
X Vi(2))
(8)
I+
Lv = Σ
Vi - ((2x)
We use an ensemble of 3 critics to reduce variance. To enable FoG optimization. is important to use a well-regularized world model. We use the (Ea(s, e), (s e), (s, a, e)) model proposed by TD-MPC2 Hansen al. (2024) with learnable task embeddings e. trained in an auto-regressive fashion by sampling data from buffer with loss function:
Lwm (Φ) E. eea.. B
Σ E>' (1 - sg( e)) CCE(fe
(10)
where sg(+) is the stop-gradient operator and CE is the cross-entropy loss function. Reward prediction formulated as a discrete regression problem in log- nog-transformed space. Furthermore, Es and Fs SimNorm activation (Eq. 4) in their output layers. All trainable models are fully-connected MLPs with LayerNorm Ba et al. (2016) and Mish activation Misra (2019). The complete algorithm is shown in Algorithm 1. Further implementation details can be found in Appendix C.
5

---

Published as a conference paper at ICLR 2025
Figure 4: High-dimensional single-task environments (left to right): Hopper. Ant, Anymal, Humanoid and SNU Humanoid. Our method successfully learns tasks with up 132 continuous action dimensions. Additional 80 multi-task environments used in this paper are listed Appendix
4 EXPERIMENTAL RESULTS
4.1 CONTACT-RICH SINGLE TASKS
The aim this section is to understand whether smooth world models create better optimization landscapes than ground-truth dynamics, facilitating efficient FoG optimization. We study this 5 continuous control tasks (Figure 4) with up to A = R 162 using the differentiable simulator dflex Xu et al. (2022). Comparisons include SHAC Xu et al. (2022), which uses ground-truth dynamics and rewards with a similar actor-critic architecture to PWM. Furthermore, we compare against two world model approaches. DreamerV3 Hafner al. (2023) learns its world model via reconstruction, its actor via ZoG optimization, and critic via Model based Value Expansion (MVE) Feinberg et al. (2018). TD-MPC2 Hansen al. (2024) uses the same world model as PWM but learns policy in model-free fashion and actively plans at inference time. We additionally include model-free baselines PPO Schulman et (2017) and SAC Haarnoja et (2018). This comparison allows us to understand whether (1) FoGibased optimization can learn better policies asymptotically and (2) whether smooth world models induce better optimization landscapes for FoG optimization.
We conduct this experiment across 5 tasks with 10 seeds each, where PWM, DreamerV3, and TD- MPC2 use pre-trained world models and are left to learn policy and fine-tune their world models online. This is done enable fair comparison SHAC, which directly has access to the simulation model and does not require any training. The results in Figure 5 reveal that (1) PWM is able to learn policies with higher reward than SHAC asymptotically, indicating that regularized world models induce smoother optimization landscapes than the true (discontinuous) dynamics. Furthermore (2), despite using the same world model, our method is able learn policies with higher rewards than TD-MPC2 without the need for online planning. However, PWM does not scale well to the highest-dimensional task. More experiment details and results are included in Appendix D.
1.00
Fraction of runs
with score ^ T
0.75
0.50
0.25
0.00
0.8
0.9
1.0
1.1
1.2
1.3
0.6
0.8
1.0
1.2
1.4
1.6
PPO-normalized Reward
Normalized Score (7)
PPO
SHAC
TD-MPC2
DreamerV3
PWM (ours)
Figure Aggregate results from high-dimensional locomotion
where each agent trained solve
just that task (i.e. specialist). The left figure summarizes rewards achieved at the end of training using 50% IQM for the solid lines and 95% CI as suggested by Agarwal et al. (2021), as well as mean for the dashed lines. We see that PWM achieves higher rewards than main baselines TD TD-MPC2 and SHAC. The right figure shows score distributions across all tasks which lets understand the performance variability of each approach.

---

Published as a conference paper at ICLR 2025
Multi-task models comparison
Multi-task vs single-task experts
MT80 MT33
50
World model training (ms)
0.7 0.8
50 Inference time (ms)
1.0
Normalized score
Normalized score
SAC
DreamerV3
TD-MPC2
PWM (ours)
30 and task set settings which include euvironments from dm control Tunyasuvunakool et al. (2020) and MetaWorld Yu et al. (2020). The results show 50% IQM with the solid lines and mean with the dashed lines. The bars represent 95% CL. In both settings PWM achieves higher reward than TD-MPC2 without the need for online planning. The middle figure compares the training and inference times of TD-MP TDMMPC2 and PWM the 48M parameter model. PWM has significantly lower inference time as it does not plan online. The right figure shows a comparison between multi-task PWM and TD-MPC2 and single-task experts SAC and DreamerV3 on the MT30 task set. Notably, PWM is able match the performance of SAC and DreamerV3.
Figure 6: Multi-task results. The ligure shows results multi-task agents
4.2 MULTI-TASK WORLD-MODEL
We analyze the scalability of our proposed framework and method large multi-task pre-trained world models. We evaluate on two settings: (1) 30 continuous control dm dm_ _om_otrool tasks Tunyasuvunakood et al. (2020) ranging from m = to m = and (2) 80 tasks, which include 50 additional manipulation tasks from MetaWorld Yu et al. (2020) with were introduced as MT30 and MT80 by Hansen et al. (2024). In conducting our experiments, we harness the same data and world model architecture TD-MPC2. The data consists of 120k and 40k trajectories per dm controll and MetaWorld task, respectively generated by 3 random seeds of TD-MPC2 runs. The world models (2024) with slight modifications to make them differentiable (Appendix C).
39 and
4. These two multi-task settings
the 48M parameter models introduced in Hansen et al.
To train PWM, first pre-train the world models on the dataset in a similar fashion to TD-MPC2, but with training H = 16 and 7 0.99 for better first-order gradients, as highlighted in Section 3.2. Then we train PWM policy on each particular task using the offline datasets for 10k gradient steps, which take 9.3 minutes on an Nvidia RTX6000 GPU. We evaluate task performance for 10 seeds for each task and aggregate results in Figure We compare against TD-MPC2, which learns multi-task policy while pre-training its world model and relies on online planning at inference. We can see that PWM learns behavior, achieving higher reward than TD-MPC2 while also being significantly faster inference time. While the fast per-task training is enabled by FoG optimization, we also find that training a single multi-task policy produces poor results, as shown in Appendix F, We further compare multi-task PWM policy to online-trainee
single-task experts SAC Haarnoja al. (2018) and Dreamer V3 Hafner et al. (2023). Figure 6 reveals that multi-task PWM, while disadvantaged, performs comparably to the single-task experts without requiring any environment interaction and only training policies for < 10 minutes per task. Additional results Appendix E.
4.3 ABLATIONS
We perform 4 ablations on the complex single task experiments in order to understand the nuances of first-order optimization through world models with PWM.
We increase the contact stiffness to be more realistic, but also more stiff contact gives gradients with high sample error Suh et al. (2022). We run the same experiment as Section 4.1, but only for the Hopper task, and present the aggregate results from 5 random seeds in Figure 7a, where normalize rewards by the maximum reward achieved by PPO in Section 4.1. We see that while PPO and PWM rewards remain similar to prior results, SHAC performance decreases by 48%. This shows that regularized world models are robust to stiff contact models and thus more generalizable than differentiable simulations.
7

---

Published as a conference paper at ICLR 2025
PPO-nom. Reward
1.0
PPO
PPO-norm. Reward
32
SHAC
512
PWM
2048
0.5
0.0
2
M
M
10
50
100
150
200
Steps
Time (m)
(a) Contact stiffinss ablation.
(b) Policy batch size ablation.
Figure 7: Left figure shows contact stiffness ablation where we increase contact stiffness on the Hopper task and analyze the effects policy learning. The results indicate that stiff (hut realistic) contact has adverse effects on SHAC which uses simulation model to learn. Meanwhile, PPO PWM remain unaffected with PWM still obtaining 17% more reward than PPO asymptotically. The right figure shows a policy batch size ablation on the Any task where we vary only the batch size used to train the policy components of PWM. Unfortunately we observe that PWM provides best result within unit of time by using small batch sizes. Both figures show 50% IQM and 95% CI over 5 random seeds.
PPO-norm. Reward
1.5
World Model Loss
Normalized Reward
0.8
TD-MPC2
PWM
1.0
0.05
0.6
ReLU
0.4
0.5
Mish
SimNorm
0.2
0.0
0.04
0
10 M
20 M
50
100 k
250 k
Steps
World model training steps (b) World model vs policy sample efficiency.
(a) World model ablation.
Figure 8: The left figure ablates the activation functions of the world model used learn policies the Ant task. We progressively add more regularization the world model via changes the activation function and observe inverse correlation between world model loss and policy reward. This indicates that we should not construct world models for accuracy but for policy learning. The right figure investigates the policy sample efficiency on 5 ddm_contte control tasks. We use the same data to pre-train world models for varying amount of gradient steps and then train the policy for 50k gradient steps and compare against TD-M DD-MPC2 (without planning). The results indicate that PWM policies are significantly more sample efficient also require better trained world models, All results shown are 50% IQM with 95% across random seeds.
The second ablation explores batch sizes for policy learning with first-order gradients. Contrary to model-free methods, which can scale large batch sizes, we find that FoG techniques like PWM benefit from smaller batch sizes. We explore this on the Ant task in Figure 7b, where we plot 50% IQM rewards over random seeds. While larger batch sizes allow us to generate more data within unit of time, that does not necessarily translate to learning better policies.
Next we ablate the world model regularization. We perform the same experiment as Section 4.1 on the Ant task but now pre-train 3 different world models. (1) with ReLU activation func., (2) with Mish activation fune. and (3) with Mish activation func. and SimNorm activation func. the output layers of Es and Fg. Figure 8a reveals that while less regularization results in lower world model eиror, that does not translate to learning better policies. Surprisingly, less regularized world models enable policies to start faster (up 1M steps) but plateau suboptimal policy. Additional results in Appendix
To understand the policy sample efficiency of PWM while controlling for the world model, perform an ablation where we pre-train train the same world model for [50k, 100k, 250k] gradient steps offline dataset. Then fix the world model and train only the policy components the
8

---

Published as a conference paper at ICLR 2025
same dataset for 50k gradient steps and measure the reward. We do this for 3 random seeds and 5 dm control tasks. We repeat the same experiment for TD-MPC2 but disable its planning component in order to understand the learning dynamics of each method's policy components. The results in Figure 8b show that the PWM policy components are significantly more sample efficient than TD-MPC2 but also require better trained world models in order to obtain high reward.
5 RELATED WORK
Reinforcement learning (RL) strategies are divided into model- and model-free approaches, with the latter not assuming a model of the environment Arulkumaran et al. (2017). Model-free approaches, such PPO Schulman et al. (2017) and SAC Haarnoja al. (2018) do not require a model of the environment and represent on-policy and off-policy methods, respectively. These algorithms use actor-critic structure, where the critic assesses the policy while the actor updates it through gradient-based optimization to maximize rewards Konda & Tsitsiklis (1999).
Gradient estimator types. In the absence of direct access to dynamics and reward functions, it is common to use the Policy Gradient Theorem Sutton et al. (1999), a zeroth-order method, to estimate gradients. Although robust to discontinuities, this method exhibits high variance, leading to sample inefficiency Mohamed et al. (2020). In contrast, first-order gradients (FoG) offer lower variance by differentiating through the objective but struggle with discontinuities Suh et al. (2022). Differentiable simulations have risen as a tool to study the properties of gradient estimators Howell et al. (2022); Metz et al. (2021) and have produced model-based algorithms that use FoG optimization through physics to learn high-performing policies Xu et al. (2022); Georgiev et al. (2024).
Multi-task models. While traditional RI. focuses on single-task policies, the broader robotics field increasingly adopting large multi-task models through behavior cloning Firoozi et al. (2023). Recent efforts like Open Padalkar et al. (2023) and Octo Octo Model Team et al. (2024) have demonstrated improved performance across various tasks and embodiments by leveraging large models and datasets. However, the potential of these large-scale approaches in RL. remains largely unexplored. While GATO Reed et al. (2022) attempted to scale model-free RL across multiple tasks, it faced challenges with sample inefficiency and required significant fine-tuning. Conversely, TD-MPC2 Hansen et al. (2024) successfully scaled a 317M parameter world model for online planning across 80 tasks. While showing impressive multi-task scalability. it failed to solve all tasks and exhibits limited scalability due to online planning. Our work builds on the world model architecture proposed by TD-MPC2 but employs FoG optimization for policy learning and extracts per-tasks policies. DreamerV3 Hafner et al. (2023) also integrates world models with FoG but focuses on online learning without addressing multi-task scenarios. Our work delves deeper into the relationship between world models and policy learning, exploring the essential characteristics of world models that facilitate efficient optimization.
6 CONCLUSION
In this study, we analyzed world models through policy gradient estimation and identified an inverse correlation between the accuracy of world models and episode rewards. We concluded that world models should prioritize smoothness and a smaller optimality gap over accuracy to enhance policy performance. Building on these insights, we propose Policy learning through Multi-task World Models (PWM), MBRL algorithm that integrates smooth world models with first-order gradient (FoG) optimization. Our evaluations showed that PWM can outperform existing methods, including those with access to ground-truth simulation dynamics, in learning high-reward policies for high- dimensional tasks. To scale to multi-task settings, propose a framework where world models are pre-trained offline and treated as differentiable simulations. Our results demonstrate that PWM can be used to learn expert policies in <10 minutes per task, achieving higher rewards without the need for expensive online planning. With ample data and large, smooth world models, we believe this approach has significant potential for scalability.
Despite its demonstrated efficacy, PWM has notable limitations. Firstly, performance relies heavily the availability of substantial pre-existing data to train the world model, which might not always be feasible. especially in novel or low-data environments. Secondly, although PWM facilitates fast and cost-effective policy training, it necessitates re-training for each new task, which could limit its applicability in scenarios requiring rapid adaptation to diverse tasks. Lastly, the current TD-MPC2 world models used are difficult to train at scale due to their nego-regressive formulation.
Limitations.

---

Published as a conference paper at ICLR 2025
made available at imgeorgiev. com/ pwm. We rely on dflex, MetaWorld, DMControl and MuJoCo for simula- tion which are publicly available under MIT and Apache 2.0 licenses. We use multi-task data from TD-MPC2 which is publicly available. Implementation details and full list of hyper-parameters are available in Appendix C.
Reproducibility statement. Code, training data and checkpoints
REFERENCES
Rishabh Agarwal, Max Schwarzer, Pablo Samuel Castro, Aaron C Courville, and Marc Bellemare. Deep reinforcement learning at the edge of the statistical precipice. Advances in neural information processing systems, 34:29304 24222,0--9930,, 2021.
Kai Arulkumaran, Mare Peter Deisenroth, Miles Brundage, and Anil Anthony Bharath. Deep reinforcement learning: A brief survey. IEEE Signal Processing Magazine, 34(6)266-38, 2017.
Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E Hinton. Layer normalization. arXiv preprint arXiv: 1607.06600 2016.
Rishi Bommasani, Drew A Hudson, Ehsan Adeli, Russ Altman, Simran Arora, Sydney von Arx, Michael S Bernstein, Jeannette Bohg, Antoine Bosselut, Emma Brunskill, et al. On the opportuni- ties and risks of foundation models. arXiv preprint arXiv:2108.07256 2021.
Konstantinos Bousmalis, Giulia Vezzani, Dushyant Rao, Coline Devin, Alex X Lee, Maria Bauza, Todor Davchev, Yuxiang Zhou, Agrim Gupta, Akhil Raju, et al. Robocat: self-improving foundation agent for robotic manipulation. arXiv preprint arXiv:2306. 2023.
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D) Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et Language models are few-shot learners. Advances in neural information processing systems, 33:1877-1901, 2020.
Vladimir Feinberg, Alvin Wan, Ion Stoica, Michael Jordan, Joseph E Gonzalez, and Sergey Levine. Model-based value expansion for efficient model-free reinforcement learning. In Proceedings of the 35th International Conference on Machine Learning (ICML 2018), 2018.
Roya Firoozi, Johnathan Tucker, Stephen Tian, Anirudha Majumdar, Jiankai Sun, Weiyu Liu, Yuke Zhu, Shuran Song, Ashish Kapoor, Karol Hausman, et al. Foundation models in robotics: Applica- tions, challenges, and the future. arXiv preprint arXiv:2312.07843, 2023.
Ignat Georgiev, Krishnan Srinivasan, Jie Xu, Eric Heiden, and Animesh Garg. Adaptive horizon actor-critic for policy learningin contact-rich differentiable simulation. In International Conference on Machine Learning. PMLR, 2024.
Ankit Goyal, Jie Xu, Yijie Guo, Valts Blukis, Yu-Wei Chao, and Dieter Fox. Ryt: Robotic view transformer for 3d object manipulation. In Conference on Robot Learning, pp. 694-710. PMLR, 2023.
Tuomas Haarnoja. Aurick Zhou, Kristian Hartikainen, George Tucker, Schoon Ha, Jie Tan, Vikash Kumar, Henry Zhu, Abhishek Gupta, Pieter Abbeel, et al. Soft actor-critic algorithms and applications. arXiv preprint arXiv: 18/2.0590 2018.
Danijar Hafner, Timothy Lillicrap, Jimmy Ba, and Mohammad Norouzi. Dream to control: Learning behaviors by latent imagination. arXiv preprint 19/2.0/0// 2019.
Danijar Hafner, Jurgis Pasukonis, Jimmy Ba, and Timothy Lillicrap. Mastering diverse domains through world models. arXiv preprint arXiv:2301.041000 2023.
Nicklas Hansen, Xiaolong Wang, and Hao Su. Temporal difference learning for model predictive control. 2022.
Nicklas Hansen, Hao Su, and Xiaolong Wang. Td-mpe2: Scalable, robust world models for continuous control. 2024.
Taylor A Howell, Simon Le CleacâÄZh, Zico Kolter, Mac Schwager, and Zachary Manchester. Dojo: A differentiable simulator for robotics. arXiv preprint arXiv:2203.0088 2022.
10

---

Published as a conference paper at ICLR 2025
Marco Hutter, Christian Gehring, Dominic Jud, Andreas Lauber, C Dario Bellicoso, Vassilios Tsounis, Jemin Hwangbo, Karen Bodie, Peter Fankhauser, Michael Bloesch, et al. Anymal-a highly mobile and dynamic quadrupedal robot. In 2016 IEEE/RSJ international conference intelligent robots and systems (IROS), pp. 38-44. IEEE, 2016.
Durk P Kingma, Tim Salimans, and Max Welling. Variational dropout and the local reparametertization trick. Advances in neural information processing systems, 28, 2015.
Alexander Kirillov, Eric Mintun, Nikhila Ravi, Hanzi Mao, Chloc Rolland, Laura Gustafson, Tete Xiao, Spencer Whitehead, Alexander C Berg, Wan-Yen Lo, et al. Segment anything. In Proceedings of the IKEE/CVF International Conference Computer Vision, pp. 4015-4026 2023.
Vijay Konda and John Tsitsiklis. Actor-critic algorithms. Advances in neural information processing systems, 12, 1999.
Aviral Kumar, Jocy Hong, Anikait Singh, and Sergey Levine. Should i run offline reinforcement learning behavioral cloning? In International conference on learning representations, 2021.
Nathan Lambert, Brandon Amos, Omry Yadan, and Roberto Calandra. Objective mismatch in model-based reinforcement learning. arXiv preprint arXiv:2002.04523, 2020.
Minji Lee, Jeongmin Lee, and Dongjun Lee. Differentiable dynamics simulation using invariant contact mapping and damped contact force. In 2023 IEEE International Conference Robotics and Automation (ICRA), pp. 1168311 1689. IEEE, 2023.
Luke Metz, C Daniel Freeman, Samuel S Scheenholz, and Tal Kachman. Gradients not all you need. arXiv preprint arXiv:2111.058033 2021.
Diganta Misra. Mish: A self regularized non-monotonic activation function. arXiv preprint
1eXi::908.0086/1 2019.
Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Alex Graves, loannis Antonoglou, Daan Wierstra, and Martin Ricdmiller. Playing atari with deep reinforcement learning. arXiv preprint 13255:13/2.5602, 2013.
Shakir Mohamed, Mihaela Rosca, Michael Figurnov, and Andriy Mnih. Monte carlo gradient estimation in machine learning. The Journal of Machine Learning Research, 21(1):5183-522 2020.
Richard M Murray and John Edmond Hauser. A case study in approximate linearization: The acrobat example. Electronics Research Laboratory, College of Engineering, University of âAç, 1991.
Octo Model Team, Dibya Ghosh, Homer Walke, Karl Pertsch, Kevin Black, Oier Mees, Sudeep Dasari, Joey Hejna, Charles Xu, Jianlan Luo. Tobias Kreiman, You Liang Tan, Lawrence Yunliang Chen, Pannag Sanketi, Quan Vuong, Ted Xiao, Dorsa Sadigh, Chelsea Finn, and Sergey Levine. Octo: An open-source generalist robot policy. In Proceedings of Robotics: Science and Systems, Delft, Netherlands, 2024.
Abhishek Padalkar, Acor Pooley, Ajinkya Jain, Alex Bewley, Alex Herzog, Alex Irpan, Alexander Khazatsky, Anant Rai, Anikait Singh, Anthony Brohan, et al. Open x-embodiment: Robotic learning datasets and rt-x models. arXiv preprint 2rXiv:2310.0880 2023.
Paavo Parmas, Takuma Seno, and Yuma Aoki. Model-based reinforcement learning with scalable composite policy gradient estimators. In International Conference Machine Learning, pp. 27346-27377. PMLR, 2023.
Scott Reed, Konrad Zolna, Emilio Parisotto, Sergio Gomez Colmenarejo, Alexander Novikov, Gabriel Barth-Maron, Mai Gimenez, Yury Sulsky, Jackie Kay, Jost Tobias Springenberg, et al. A generalist agent. arXiv preprint arXiv:2205.06//5. 2022.
John Schulman, Nicolas Heess, Theophane Weber, and Pieter Abbeel. Gradient estimation using stochastic computation graphs. Advances in neural information processing systems, 28, 2015.

---

Published as a conference paper at ICLR 2025
John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. arXiv preprint arXiv:1707.06347, 2017.
Hyung Ju Suh, Max Simchowitz, Kaiqing Zhang, and Russ Tedrake. Do differentiable simulators give better policy gradients? In International Conference Machine Learning, pp. 20668-20696. PMLR, 2022.
Richard S Sutton and Andrew G Barto. Reinforcement learning: An introduction. MIT press, 2018.
Richard S Sutton, David McAllester, Satinder Singh, and Yishay Mansour. Policy gradient methods for reinforcement learning with function approximation. Advances in neural information processing systems, 12, 1999.
Saran Tunyasuvunakool, Alistair Muldal, Yotam Doron, Siqi Liu, Steven Bohez, Josh Merel, Tom Erez, Timothy Lillicrap, Nicolas Heess, and Yuval Tassa. dm__cmttro Software and tasks for continuous control. Software Impacts, 6:100022, 2020.
Jie Xu, Viktor Makoviychuk, Yashraj Narang, Fabio Ramos, Wojeiech Matusik, Animesh Garg, and Miles Macklin. Accelerated policy learning with parallel differentiable simulation. In International Conference on Learning Representations, 2022.
Tianhe Yu, Deirdre Quillen, Zhanpeng He, Ryan Julian, Karol Hausman, Chelsea Finn, and Sergey Levine. Meta Merl-world: A benchmark and evaluation for multi-task and meta reinforcement learning. In Conference on robot learning, pp. 1094-1100 110. PMLR, 2020.
Brianna Zitkovich, Tianhe Yu, Sichun Xu, Peng Xu, Ted Xiao, Fei Xia, Jialin Wu, Paul Wohlhart, Stefan Welker, Ayzaan Wahid, Quan Vuong, Vincent Vanhoucke, Huong Tran, Radu Soricut, Anikait Singh, Jaspiar Singh, Pierre Sermanet, Pannag R. Sanketi, Grecia Salazar, Michael S. Ryoo, Krista Reymann, Kanishka Rao, Karl Pertsch, Igor Mordatch, Henryk Michalewski, Yao Lu, Sergey Levine, Lisa Lee, Tsang- Wei Edward Lee, Isabel Leal, Yuheng Kuang, Dmitry Kalashnikov, Ryan Julian, Nikhil Joshi, Alex Irpan, Brian Ichter, Jasmine Hsu, Alexander Herzog, Karol Hausman, Keerthana Gopalakrishnan, Chuyuan Fu, Pete Florence, Chelsea Finn, Kumar Avinava Dubey, Danny Driess, Tianli Ding, Krzysztof Marcin Choromanski, Xi Chen, Yevgen Chebotar, Justice Carbajal, Noah Brown, Anthony Brohan, Montserrat Gonzalez Arenas, and Kehang Han. Rt-2: Vision- www---ae-aaattinnnn models transfer web knowledge to robotic control. In Jie Tan, Marc Toussaint, and Kourosh Darvish (eds.), Proceedings of The 7th Conference Robot Learning, volume 229 of Proceedings of Machine Learning Research, pp. 2165-2183, PMLR, 06-09 Nov 2023. URL https: http://poooccccccccccooooooo nggs mlr, htt.:.ppeesss./e/eee http:///w////tttoiioooo/e html.
12

---

Published as a conference paper at ICLR 2025
20
ReLU
ReLU SimNorm
SimNorm
100
0
50
-10
250
5000
750
1000
-2
100
Iteration
e
Iterations
Figure 10: Extended ball-wall toy example results. The left figure shows the model losses as they are trained to approximate the target function J(0). The middle figure shows the output of the trained model across the spectrum of well as the true target. This is the function attempt to minimize in the right figure. We see that when using the MLP with ReLU activation functions, the optimizer quickly gets stuck local minima while the model using SimNorm activation function is able to find a solution closer to the true one.
A BALL-WALL EXAMPLE DETAILS
This section provides more details on the ball-wall example used to showcase the issues of optimizing through contact in Section 3.1. constructing this toy example we chose a simple physical system that exhibits contact discontinuities. Inspired by Suh al. (Suh et al., 2022). we constructed simple problem of a point mass (ball) being thrown forward (x direction) at a fixed velocity v. The optimization parameter of interest is the initial angle 0 and the goal is maximize forward distance traveled (in 2D). For simplicity we assume that the ball sticks to the wall (without complex contact) which can be expressed as:
J(0) = F1 = f(0) = To + e cos + 2 gt if Veontact >> h
(11)
else
where g 9.81 is gravity, h and w are the height and width of the wall, (ro, yo) is the starting position, v 10 the starting velocity and is time. Veontact the height the time of contact tcontact which are both given by solving Eq. 11 for f(0) w:
feontact -C)o((0)
v cos (0) awe
Yeuntact = yo + sin( +
2 gt2
We visualize the toy example in Figure aid reading. With (0) defined, we attempt to learn it with two Multi-Layer Per- ceptrons (MLPs). We configure them to have 2 hidden layers of 32 neurons each. The first MLP uses ReLU activation func- tions, while the latter uses SimNorm activation functions as defined in Eq. 4. Both models are initialized with identical random parameters and are trained with the ADAM optimizer with learning rate = 2 10 for 100 epochs using batch size of B 50. The data we use to train the models was 1000 uniform samples of f0) within 9 [-Л, r]. Figure 10 shows the training losses of the models, induced optimization land- scapes and the losses when attempting to maximize the models as you would do in an RL setting.
h
y
Figure 9: Pedagogical ball-wall toy problem visualized.
13

---

Published as a conference paper at ICLR 2025
10-2
104
Var[v/(O)]
ESNR(V)(())
10-7
True
10
True
99
10
MLP H=3
10-2 -
MLP H=3
MLP H=16
MLP H=16
S
0 20
40
0 20
40
H
Figure 11: Double pendulum pedagogical example. The middle figure evaluates the variance of policy gradient estimates over N = 100 Monte- Carlo samples for varying horizons H. The right figure shows the same data but plots the Expected Signal-to-Ner -o-isee ratio (ESNR) with higher values translating These results suggests that world models trained over long horizon trajectories provide more useful gradients.
useful gradients.
B DOUBLE PENDULUM EXAMPLE DETAILS
The double pendulum (also known Acrobot (Murray &
Hauser, 1991)) is a classic under-actuated chaotic system. It is characterized by its sensitivity to initial conditions where even small perturbations result in large gradient variance with long horizon (> 20) trajectories. We chose this system to analyze variance and expected signal-to-noise -o-sse ratio (ENSR) in Section 3.2 it is the easiest problem exhibiting chaosness. We model this toy problem similar to DMControl (Tunyasuvunakoo) et al., 2020) in our differentiable simulator, dflex as we need ground truth gradients for comparison. The first link with angle is fixed the base and not actuated. The second link with angle 02 is the only control input via 0g. The state of the system id calculated as:
[cos(O) sin(01), (0s) sin s02), 01, 0z]
=
The objective of this toy example is to bring and balance the pendulum upwards which achieve by formulating a reward:
r(s, a) - 0.16
Next we train world models to approximate the dynamics and reward above. For this we collect data with the SHAC algorithm (Xu et al., 2022) over different runs for total of 24,000 episodes of length 240 timesteps. Maximum episode reward achieved during data collection was -942.95. Then we train two TDMPC2 (Hansen et al., 2024) world models on the collected data. We use the 5M parameter model which features a latent state of dimension of 512, encoder Es with one hidden dimension of 256, dynamics model MLP with 2 hidden layers with 512 neurons and rewards model of the same design. We keep the same hyper-parameters as per the origin work by Hensen et al. but use 7 0.99 which we found reduce variance substantially. We train two models with different training horizons H = 3 and H = 16 for 100k batch samples and batch size of 1024.
With the trained models, we now compare the variance of stochastic gradients provided by the true dynamics of the simulation and the two trained models. We do this by loading the best policy learned by SHAC during data collection and executing H/ 50 rollout across the models. We ensure that the same actions are taken for each evaluated models and collect 100 Monte Carlo samples. In addition to variance, we report ESNR suggested by (Parmas al., 2023) and defined in 5. Higher ESNR translate to more useful gradients and we naturally should expect values to decrease with increased H. We reported the results in Figure but also duplicate them Figure 11 for convenience and ease of reading.
14

---

Published as a conference paper at ICLR 2025
C IMPLEMENTATION DETAILS AND HYPER-PARAMETT
The section details several implementation details of PWM that we thought are not crucial for understanding the proposed approach in Section 3.3 but are important for replicating the results.
1. Reward binning the reward model we use in PWM is formulated as a discrete regression problem where R rewards are discretized into a predefined number of bins. Similar to (Hansen al., 2024; Lee al., 2023), we this enable robustness to reward scale and multi-task-ness. In particular, we perform two-hot encoding using SymLog and SymExp operators which are mathematically defined as:
SymLog Two-hot encoding is then performed with:
(x) ) = sign(x) log(1 + |2r))
SymExp(a) = sign(x) ( 1)
def twe_hot (x)
clamp (symlog (x), vmin, vmax)
bin_idx #loor (x vmin) bin_size) bin_offaet (x vmin) bin_size bin_idx
soft_twe_hot zeros (x.size (0), num_bins)
soft_two_hon [bin_idx] - bin_offset soft_two hot [bin_odx odx return swot_ hot
sin_offfet
Inverting this operation to get back to scalar rewards would usually involve SymFxp(1) but note that the sign(x) operator is not differentiable and would therefore not work for FoG. Instead, we chose to omit the SymExp(1) operation which technically now returns pscudo-rewards but also gradients which we found sufficient for policy learning:
def two jot coot_inersion (x)
vals linspace (vmin, vmax, bins)
softmax
x torch. sum vals, dim=-1)
return x
2. Critic training while Algorithm function to similar results as presented in 4, we found it beneficial to split the critic training data from single rollout into several smaller mini- batches and over them for multiple gradient steps. In our implementation we split the data into 4 mini-batches and perform 8 gradient steps over them with uniform sampling. With 16 and batch size 64, this translates to critic batch size of 256.
3. Minimum policy noise- Due to the larger amount of gradient steps needed, we noticed that PWM's actor tends to collapse to deterministic policy rapidly. As such, we found beneficial to include lower bound on the standard deviation of the action distribution in order to maintain stochasticity in the optimization process. We have used 0.24 throughout this paper. While similar results would be possible by adding an entropy term (Schulman et al., 2017), we found our current solution sufficient
4. World model fine-tuning Throughout all of our experiments we found that the offline data used to train PWM's world model be crucial to learning a good policy. In very nigh-dimensional tasks such as Humanoid SNU, collecting extensive data is difficult task. As such, in these tasks we found beneficial to online fine-tune the world model. We do this on all single-task experiments of Section 4.1 using the default hyper-paraments parameters and replay buffer of size 1024.
15

---

Published as a conference paper at ICLR 2025
Hyper-parameter
Value
Policy components
Horizon (H)
16
Batch size
64
0g
x 10
Oleb
10-4
Actor grad norm
Critic grad norm
100
Actor hidden layers
4400, 200, 100]
Critic hidden layers
400,2000
Number of critics
3
0.95
Y
0.99
Critic batch split
Critic iterations
World model components (48M)
Latent state (z) dimension
768
Horizon (H)
16
Batch size
1024
Ol
3 x 10
World model grad norm
20.0
SimNorm V
8
Reward bins
101
Encoder E& hidden layers
[1792, 1792, 1792]
Dynamics F& hidden layers
[1792, 1792
Reward R& hidden layers
(1792, 1792
Task encoding dimension
96
Table 1: Table of paraeeparameters used in PWM, shared across all tasks.
Figure 12: Locomotion environments (left right): Hopper, Ant, Anymal, Humanoid and SNU Humanoid.
D CONTACT-RICH SINGLE TASK EXPERIMENT DETAILS
In Section 4.1, we explore locomotion tasks with increasing complexity. They are described below and shown in Figure 4.
1. Hopper, a single-legged robot jumping only in one axis with n = 11 and m 3. Ant, a four-legged robot with n = 37 and m = 8.
Anymal, a more sophisticated quadruped with n = 49 and m 12 modeled after (Hutter et al., 2016).
Humanoid, a classic contact-rich environment with 76 and exploration to find good policy.
21 which requires extensive
5. SNU Humanoid, a version of Humanoid lower body where instead of joint torque control, the robot is controlled via 152 muscles intended to challenge the scaling capabilities of algorithms.
All tasks share the same common main objective maximize forward velocity uz:
16

---

Published as a conference paper at ICLR 2025
Hopper
Ant
Anymal
Humanoid
SNU Humanoid
Ep reward
5 M
10 M.
M M
M Si Simulation steps
20 Simulation steps
10
10 20 M
Simulation steps
Simulation steps
Simulation steps
TD-MPC2
DreamerV3
PWM (ours)
Figure 13: Learning curves for each environment. This figure shows 50% IQM and 95% CI across 10 random seeds for each task in the dflex simulation suite. Rewards are normalized by the maximum reward achieved by PPO (usually 100 steps). While PWM remains competitive with SHAC for most tasks, does scale well the 152 action dimension SNU Humanoid. Note that SHAC uses gradients from the differentiable simulation directly and converges in fewer samples, explaining the truncated curves.
Environment
Reward
Hopper
Us
Rhoight Rangle 0.1 aal
Ant
Ur
Rheight + 0. A.lBangle + Rheading 0.01
Anymal
Up + Rheight + 0. Rangle + Rheadsng 0.01
Humanoid
Up Rheight + 0. Rangle Rheadinn 0.002
Humanoid STU
Ur Rheight + 0. Rangle Rheading 0.002 a
Table 2: Rewards used each task bench maach-marked in Section
We additionally use auxiliary rewards Rheight to incentivise the agent to, Rangte to keep the agents' normal vector point up, Rheading keep the agent's heading pointing towards the direction of running and a norm over the actions to incentivise energy- energy-efficient policies. For most algorithms, none of these rewards apart from the last one are crucial to succeed in the task. However, all of them aid learning policies faster.
h hterm
ifh 2 htrr
Rhoight
-200 hearm) )
ifh < heerm
Rangle =
Overm
Rangle = [[forworrd Qagent] 2 is the difference between the heading of the agent Jagent and the forward vector Qagent- h is the height of the CoM of the agent and is the angle of its normal vector. herm and ftem are parameters that we set for each environment depending on the robot morphology. Similar other high-perffrmance RL applications simulation, we find crucial episode early if the agent exceeds these termination parameters. However, it is worth noting that AHAC is still capable of solving all tasks described in the paper without these termination conditions, albeit slower.
terminate
All results presented in Figure 13 are for 10 random seeds using the simulator in vectorized fashion with 64 parallel environments for all approaches, except PPO which uses 1024. We note that while simulation steps appear high, of these experiments are executed hours on an Nvidia RTX6000 GPU. In addition the the learning curves of Figure 13, also present tabular results below:
We note that TDMPC2 and PWM use pre-trained world models on 20480 episodes of each task. The world models are trained for 100k gradient steps and the same world models (specific to each task) are loaded into both approaches. The data consists of trajectories of varying policy quality generated with the SHAC algorithm. Trajectories include near-O episode rewards well as the highest reward achieved by SHAC. Note that we also run an early termination mechanism in these tasks which is done to accelerate learning and iteration.
17

---

Published as a conference paper at ICLR 2025
Hopper
Ant
Anymal
Humanoid
SNU Humanoid
PPO
1.00 + 0.11
1.00 + 0.12
1.00 0.03
1.00 0.05
1.00 0.09
SAC
0.87 + 0.16
0.95 + 0.08
0.98 0.06
1.04 0.04
0.88 0.11
DreamerV3
1.15 + 0.47
1.12 + 0.45
1.18
0.47
1.03 0.43
0.48 = 0.21
TDMPC2
0.85 + 0.37
1.07 + 0.44
0.98 + 0.48
1.05 + 0.46
0.26 = 0.12
SHAC
1.02 + 0.03
1.16 + 0.13
1.26 0.04
1.15 + 0.04
1.44 0.08
PWM
1.20 = 0.29
1.46 + 0.31
1.16 + 0.24
1.19 = 0.025
1.36 = 0.56
Table 3: Tabular results of the asymptotic (end of training) rewards achieved by each algorithm all tasks. results presented are PPO-normalised IQM and standard deviation across 10 random seeds. Most. algorithms have been trained until convergence.
Hopper
Ant
Anymal
Humanoid
SNU Humanoid
PPO
4742 + 521
6605 + 793
12029 360
7293 365
4114 + 370 3620 + 453
SAC
4126 + 759
6275 + 528
11788 + 722
7285 292
DreamerV3
5436 + 2213
7370 3002
14169 = 5692
7546 + 3106
2018 + 873
TDMPC2
4027 + 1768
7080 2885
11787 = 4702
7634 3317
1121 - 525
SHAC
4837 + 142
7662 + 859
15157 + 481
8387 = 292
5924 + 329 5767 + 2394
PWM
5680 + 2303
9672 2012
13927 2882
8661 1792
Table 4: Tabular results of the asymptotic (end of training) rewards achieved by each algorithm across all tasks. The results presented are 50 % IQM and standard deviation across 10 random seeds. All algorithms have been trained until convergence.
E MULTI-TASK EXPERIMENTS ADDITIONAL RESULTS
In this section we provide additional results on multi-task experiments. While we find it beneficial train the world model the same horizon as the policy learning, it is not strictly necessary to achieve good performance. In Figure 14 we present an ablation where we compare PWM world models pre-trained results reveal that H = 16 trained world models have only marginally higher scores. On deeper inspection, most of increased scores come form dm control tasks which are harder than MetaWorld tasks average. Therefore if training new world models, we advise using higher H; however if other pre-trained world models exist with suboptimal H, they will probably be also useful.
horizons
3 and H 16 and policies trained only with H 16. These
MT
MT80
0.6
0.7
0.8
0.9
Normalized score
TD-MPC2
PWM H=3
PWM H=16
Figure 14: Horizon ablation the world model
Figures 15 and 16 give scores for individual tasks for TDMPC2 and PWM across both the MT30 and MT80 task sets. We can observe that most of the increased performance of PWM is in dm__ontrol tasks which are on average more difficult than MetaWorld.
18

---

Published as a conference paper at ICLR 2025
TD-MPC2
acrobot-swingup
PWM
cartpole-balance
cartpole-balan
cartpole-swingup
cartpole-swinnuu-sparre
cheetah-jump
cheetah-run
cheetah-run-back
cheelah-run-baccwards
cheetah-run-foom
cup-catch
cup-spin
finger-spin
finger-tun-easy
finger-tum-hard
fish-swim
hopper-hop
hopper-hop-backwards
hopper-stand
pendulum-spin
pendulum-swingup
reacher-easy
reacher-hard
reacher-threeccasy
reacher-three-han
walker-run
walker-run-backkwadds
walker-stand
walker-walk
walker-walk- bakkkwrwall
0.0
0.2
0.4
0.6
0.8
1.0
Normalized score
Figure 15: Individual task results for MT30 task set.
19

---

Published as a conference paper at ICLR 2025
TD-MPC2
TDMPC2
acrobot-swingup
PWM
mw-hammer
PWM
carlpole-balance
mw-hand-insert
carppol-balannce-paareo
mw-handle-press
cartpale-swingup
mw-handle-press-
cartpolo-swingueee
serpoll-swinguu--spass
mw-hand elpuull
nw-handie-pun -aldddd---lll-d
chestat-jump
cheetah-run
nw-lever-puil
cheetah-run-back
mw-peg-insert-se pee---ittt--s
nw-peg-umplug-seee
chectah-run-fromt
mww-pick-outee
cup-catch
mw-pick-place
mw-pick-place-wit
cup-spin
finger-spin
mw-plate-slide
finger-tu--easy
mw-plate
ale-slide-baad
finger-tum-hard
fish-swim
mnw-plate-slide-aie
hopper-hop
mw-push
hopper-hop-acckwards
nw-push-back
hopper-stand
mw-push-wall
mw-assembly
baskelballl
mw-reach-wall
mw
mw-bin-picking
mw-shelf-place
mw-box-close
mw-socoar
mw-button-press
mw-stick-pull
mmw-button-press-topaaae
mw-stick--pstt
mw-button-press
mw-sweop
mw-sweep-into
mw-coffee-buuton
mw-window-close
mw-coffee-pull
mw-window-open
mw-coffee-push
pendulum--pinn
mw-dial-tum
sendulum-swingup
mw-disassemble
reacher-easy
mw-door-close
reacher-hard
mw-door-lack
reacher-
Bhreher-rree-easy
mw-docr-open
reacher-three-hard
mw-door-unlock
walker-run
mw-drawer-cdee
walker
backwaaddds
mw-drawar-open
walker-stand
nw-faucet-closs
walker-walk
mw- faucet-pep-
walker-walk-bacckward
0.00
0.25
0.50
0.75
1.00
0.00
0.25
0.50
0.75
1.00
Normalized score
Normalized score
Figure 16: Individual task results for MT80 task set.
F ADDITIONAL ABLATION RESULTS
F.1 WORLD MODEL REGULARIZATION ABLATION
We extend our world model regularization ablation from Section 4.3 to 3 additional dflex locomotion tasks Hopper, Anymal and Humanoid. The complete results are presented in Figure 17 and follow the same format as the original results in Figure 8a 50% IQM with 95% CI across 3 seeds. From the results we can observe that progressively adding more regularization, results in higher policy rewards. This follows the ablation results in our paper and the one of the core contribution of work more accurate world models, not results better rewards. Instead building world models for accuracy, we should build them to enable better policy optimization and thus higher rewards.
F.2 TRAINING A MULTI-TASK POLICY
One of the contributions of our work is the multi-task learning framework of PWM. Instead of training single multi-task policy, we propose training single multi-task world model with supervised objective and then extracting RL policies from it very efficiently. In this section justify this empirically by reproducing the MT30 experiment of Section 4.2 with two additional baselines TD-MPC2 without planning and PWM with single multi-task policy. This strips down both algorithms to very similar, the major difference being that TD-MPC2 features an off policy SAC-like policy while PWM features on-policy SHAC-like policy. Both are fed the one-hot
20

---

Published as a conference paper at ICLR 2025
Hopper
Anymal
Humanoid
PPO-norm. Reward
1.5
1.0
0.5
0.0
M
10 M
5 M Steps
10
15 M
5 M
10 M
15 M
5 M 10 M 15 Steps
Steps
Steps
ReLU
Mish
SimNorm
Figure 17: World model regularization ablation. This figure extends the ablation results of Figure 8a with additional tasks. The figure shows 50 % IQM with 95% C1 over 5 seeds per task.
Hopper
Anymal
PPO-norm. Reward
1.5
1.0
0.5
0.0
2.5
7.5 M Steps
10 M
10M
30 M
Steps
Steps
TD(A)
PWM
Figure 19: TD(A) ablation. This figure compares vanilla PWM with a version that uses a TD(A) actor objective. Results shown are 50% IQM with 95% CI over seeds.
task embeddings. The results in 18 show that both fail solving MT30 benchmark. In the case of TD-MPC2, this reveals its reliance planning and in the case of PWM, this reveals the need to train policies per-task. We believe this is due the unstable optimization objective of the RL problem.
F.3
TD(A) ABLATION
Here we replace the actor objective of PWM with TD(A) similar Dreamer (Hafner et al., 2019).
Multi-task MT30
We evaluate this on 3 single-task experiments Hopper, Ant and Anymal and show the results 19. We can that PWM overall achieves higher rewards, which we believe is to more simple and less noisy gradients obtained via TD(N). Addi- tionally, we found that TD(A) uses approx. 10% more computation due to having compute more gradients. However, it is worth noting that TD(A) learning curves appear more stable with higher dimensional tasks such as Anymal and we believe this ablation requires more large-scale studying.
0.2
0.4
0.8
Normalized
TD-MPC2 (no planning)
PWM saigle-paliey)
TD-MPC2
PWM
Figure 18: Framework ablation ablation. Extends the results of Figure 6 with two additional baselines: TD-MPC2 without planning and PWM with single multi-tae task policy learned figure shows 50% IQM results with solid lines, mean with dashed lines and 95 % CI.
all embeddings. The
21


Reviews:
[{"review_id": "JoMiInk291", "reviewer": "Reviewer_K87C", "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.", "questions": "Please see the weaknesses section.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. This paper is clearly written and easy to follow. \n2. This paper presents sufficient experimental results to demonstrate the validity of its proposed method.", "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.\n2. 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?\n3. 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.\n4. 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.\"\n5. In Lines 348-349, the mean of 'm' and 'n' are not explicitly defined or mentioned.", "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": "KgCPF4E6rJ", "reviewer": "Reviewer_WYHT", "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.", "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.\n\n2) The use of V in Eqn 4 is confusing, because I don’t think it has any relation to Eqn 2 (or does it?)\n\n3) 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? \n\n4) 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?\n\n5) 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?\n\n7) 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. \n\n8) Is performance sensitive to the simnorm hyper parameters?\n\n\nMinor: \n- 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?\n- 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.\n- 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’?\n- Typo in caption of Fig 1 (‘prpose’)\n- Typo line 188 (‘optimality hap’)\n- Typo line 377 (‘generalizble’)", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 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. \n\n- Good ablation experiments investigating sensitivity to the environment dynamics and the degree of world model regularization.", "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. \n\n- 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.", "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": "ocSWH8mxsL", "reviewer": "Reviewer_d6vp", "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.", "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?\n2. As mentioned in weakness 2, what will you consider to be the major differences between your method and MAAC? \n3. 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.\n4. 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.\n5. 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?\n6. Some curves in Figure 13 are incomplete. What's the reason for that?\n\nI am happy to increase my score if the authors justify these questions.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "1. The key idea of the paper is well delivered using pedagogical examples and is easy to follow\n\n2. The results are highly reproducible with detailed experimental settings and code included, which also benefits the community for further study.\n\n3. The experiments conducted in this paper are quite comprehensive and solid.", "weaknesses": "1. Typos and incorrect format\n- 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.\n- Typos in Equation (2): it should be $\\gamma^{h-t+1}$ instead of $\\gamma^h$\n- Typos in Algorithm 1: \"$\\nabla$\" is omitted in all equations regarding learnable parameters update.\n- Typos in Figure 10: missing \")\" in the caption.\n\n2. Lack of novelty\n- 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.\n\n3. A total training time cost comparison between PWM and other algorithms is not revealed in the paper.\n\n[1] Clavera, Ignasi, Yao Fu, and Pieter Abbeel. \"Model-Augmented Actor-Critic: Backpropagating through Paths.\" International Conference on Learning Representations (2020).", "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": "C5qTFOmwpP", "reviewer": "Reviewer_KWfz", "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.", "questions": "Besides, there are several minor questions.\n\n1. 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?\n2. 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.\n3. In Figure 6 (left), why PWM significantly underperforms TD-MPC2?\n4. In Line 357, does the baseline TD-MPC2 use the same hyperparameter as $H=16$?\n5. Can the author explain why PWM benefits from smaller batch sizes?\n6. Typo: Line 036 prpose => propose\n\nOverall, I find this paper interesting and informative. If the authors address my concerns well, I am willing to raise my rating.", "limitations": "", "rating": 6, "confidence": 5, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "1. Didactic examples are informative.\n2. Experiment results: sample efficiency to learn a task within 10 minutes seems promising.\n3. The paper is well-written and easy to understand.", "weaknesses": "The major concerns are novelty and clarity of contributions. This paper does not well highlight its own contribution to the existing literature.\n\n- 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.\n- 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.", "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)