Task-aware world model learning with meta weighting via bi-level optimization
Authors: Huining Yuan
Hongkun Dou
Xingyu Jiang
Yue Deng
Venue: NeurIPS 2023 poster (2023)
Keywords: Model-based reinforcement learning
world model
generative model
meta-learning
bi-level optimization
TL;DR: A task-aware framework for world model learning in model-based reinforcement learning.

Abstract: Aligning the world model with the environment for the agent’s specific task is crucial in model-based reinforcement learning. While value-equivalent models may achieve better task awareness than maximum-likelihood models, they sacrifice a large amount of semantic information and face implementation issues. To combine the benefits of both types of models, we propose Task-aware Environment Modeling Pipeline with bi-level Optimization (TEMPO), a bi-level model learning framework that introduces an additional level of optimization on top of a maximum-likelihood model by incorporating a meta weighter network that weights each training sample. The meta weighter in the upper level learns to generate novel sample weights by minimizing a proposed task-aware model loss. The model in the lower level focuses on important samples while maintaining rich semantic information in state representations. We evaluate TEMPO on a variety of continuous and discrete control tasks from the DeepMind Control Suite and Atari video games. Our results demonstrate that TEMPO achieves state-of-the-art performance regarding asymptotic performance, training stability, and convergence speed.

Task-aware world model learning with meta weighting via bi-level optimization
Huining Yuan
Hongkun Dou School of Astronautics, Beihang University, Beijing, China (hnyuan, douhk, jxy33zrhd, ydeng)@buaa. edu. en
Xingyu Jiang Yue Deng"
Abstract
Aligning the world model with the environment for the agent's specific task is crucial in model-based reinforcement learning. While value-equivalent models may achieve better task awareness than maximum-llkkeliood models, they sacrifice a large amount of semantic information and face implementation issues. To combine the benefits of both types of models, we propose Task-aware Environment Modeling Pipeline with bi-level Optimization (TEMPO), a bi-level model learning framework that introduces an additional level of optimization on top of a maximum-likelihoo model by incorporating meta weighter network that weights each training sample. The meta weighter the upper level Iearns to generate novel sample weights by minimizing a proposed task-aware model loss. The model in the lower level focuses on important samples while maintaining rich semantic information in state variety of continuous and discrete representations. We evaluate TEMPO Suite and Atari video games. Our control tasks from the DeepMind Control results demonstrate that TEMPO achieves state-of-the-art performance regarding asymptotic performance, training stability, and convergence speed.
Introduction
Reinforcement learning (RL) achieves intelligent behavior by optimizing sequential decision-making through a trial-and-error process (Sutton and Barto, 2018). While RL has shown outstanding success in tasks like Go and video games, the enormous quantity of samples required to train such agents poses great limitations on RL's application in real-world scenarios involving human operators, real robots, computationally expensive simulators (Moerland al., 2023).
Model-based reinforcement learning (MBRL) aims to enhance the sample efficiency and generaliza- tion capability of RL agents through two interleaved stages: model learning and behavior learning. In the model learning stage, an approximate world model of the environment is learned using real environmental samples provide the agent with the ability to generate simulated experiences pre- dict the outcome of actions. In the behavior learning stage, the agent learns its policy by interacting with the model without having take actions in the real environment. This paradigm has received lot of attention, and significant progress has been made in both model learning and behavior learning. (Sutton, 1991; Ha and Schmidhuber, 2018; Janner al., 2019; Kaiser al., 2019; Hafner al., 2019a; Schrittwieser et al., 2020).
One of the most common approaches to building model is to learn deep generative model through maximum likelihood estimation (MLE) on environmental trajectories (Buesing et al., 2018; Hafner al., 2019b,a, 2020; Ozair et al., 2021). Such models can leverage advances in probabilistic modeling, and can easily be combined with state-of-the-an model-free RL agents for task behavior (Ha and Schmidhuber, 2018; Kaiser et al., 2019). However, MLLE reconstructs all information from
"Corresponiing author
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

environmental observations equally, overlooking the information needed for learning specific task behavior. Since the model is only an approximation of the real environment, model errors can create a gap between maximizing model return and maximizing environment return, leading to poorly learned policies (Lambert et al., 2020). Therefore, it desirable to minimize such a gap by learning a model that can accurately predict those states that have a higher impact on the agent's ' task policy.
Driven by this goal, value equivalent models are designed predict future state (or state-action) values rather than the raw observations, such that the model learns to preserve only value-relevant characteristics of the environment (Farahmand et al., 2017; Schrittwieser al., 2020; Zhang et al., 2020; Grimm et al., 2020, 2021; Antonoglou et al., 2021: Nikishin et al., 2022). Nonetheless, such models often rely on specially-tailore et al., 2020), and may face challenges with implementation and optimization (Farahmand et al., 2017), constraining their scalability and robustness across different tasks and learning strategies. Additionally, the substantial amount of semantic information that is discarded during such value- focused learning can be useful for learning an effective policy.
planning algorithms for behavior learning (Schrittwieser
Is there a favorable trade-off between MLE-based models and value equivalent models, where we can enjoy the merits of both worlds? In this work, inspired by recent advances in meta-learning (Nichol et al., 2018; Shu et al., 2019; Jiang et al., 2022), we propose a bi-level framework for model learning, in which we introduce an additional level of optimization on top of an MLE-based model by incorporating a meta weighter network that assigns importance weights to each training sample in the MLE objective function. The meta weighter is then trained to generate novel sample weights by minimizing proposed task-aware model loss. Under this hierarchical framework, the meta weighter in the upper level of optimization learns to prioritize those samples with positive impact on closing the task-relevant gap between the environment and model. The model in the lower level is then forced to focus on important samples, while still learning to reconstruct environmental observations and thus, form state representation with rich semantic information to facilitate policy learning. We name our framework Task-aware Environment Modeling Pipeline with bi-level Optimization (TEMPO).
We build our bi-level framework TEMPO on top of Dreamer (Hafner et al., 2020), a powerful MBRL algorithm with a sequential VAE-like model and an actor-critic agent. A simple approximation of the gradient is proposed to update the meta weighter efficiently during implementation. We evaluate the novelty of TEMPO on challenging visual-based RI. benchmarks, including both continuous control tasks from the DeepMind Control Suite (Tassa et al., 2018) and discrete control tasks from Atari video games (Bellemare et al., 2013). Results show that our framework achieves state-of-thh-an performance compared with model-free RL algorithms, and exceeds the original DreamerV2 in terms of asymptotic performance, training stability, and convergence speed. Furthermore, we perform ablation studies to demonsirate the advantage of our proposed meta-weighting mechanism.
2 Task- Environment Modeling Pipeline with bi-level Optimization
In this work, we focus on visual-based RL control tasks, which are commonly formulated as partially observable Markov decision processes (POMDPs) with discrete time steps E [1 T]. high- dimensional observations 01.7 (usually images in visual-based cases), continuous or discrete vector- valued actions 11.7, and scalar rewards TT.7. The observations and rewards are generated by the black- box environment o,, F2 ~ p(o,, relocs, ), with actions generated by the agent ag ~ p(ay aze). The goal of RI. agents is to maximize the expected sum of rewards F.,, (E rt).
In this section, we first briefly introduce the world model from DreamerV2 (Hafner et al., 2020) as the foundation of our work. Then, we propose a loss function for evaluating the task awareness of such MLE-based world model. Finally, we introduce our main contribution, a bi-level model learning framework where a meta- weighteig mechanism is proposed to subtly combine the world model with our task-aware model loss.
2.1 The world model
We start with the Recurrent State-Space Model (RSSM) proposed in DreamerV2 (Hafner et al., 2020). The RSSM is probabilistic graphical model similar to a sequential VAE (Kingma and Welling, 2013; Sohn et al., 2015). conditions on past observations and actions to model the distribution of state transitions that have occurred in the environment.

---

Infer states, values, meta weights: steps 3-6
(a)
Data
6
States
Values
N
Meta weights
Lower optimization world model: step
Netwoks paramelers
(b)
e lower
|00
0
Inputs & outputs
Workflow of inference
Upper optimization meta weighter: step
Workflow lower optmization
(c)
a
aL 1pp4 /00
X
2² L Juwer /00 dcx
α
Workflow of upper optimization
Figure 1: Main workflow of the proposed bi-level model learning paradigm TEMPO, corresponding to steps 3-8 Algorithm 1. We start with (a) inferring the states and meta weights with model Mg. value function V, and meta weighter No. then sequentially perform (b) the lower optimization the model and (c) the upper optimization on the meta weighter.
Consider trajectory in the environment, which comprises sequence of observations, actions, and rewards for, 4s, r. The RSSM introduces deterministic states hırr and stochastic states $1.7 latent variables for each time step. The observations and rewards are then conditionally generated from these states. Specifically, the RSSM consists of 4 main components
Deterministic state module:
h1 fa 81-1, at-1)
Stochastic state module:
St ~ (s(lt
Observation & reward module:
or, po po( rel 8t)
Representation module:
8t ~ go
The deterministic state module recurrently outputs hG whereas the stochastic state module, which corresponds to the prior network in conditional VAE, predicts the prior at each time step. These two state modules work together to perform the state transitions, while the observation & reward module acts decoder reconstruct o, and r. from hn and 8t. To enable end-to-end training through variational inference (Kingma and Welling, 2013), an additional representation module, which serves as an encoder, is used to infer the posterior st. For clarity, we denote the entire model together as Mg with parameters 0.
All components of the model are trained to maximize variational bound on the trajectory log- likelihood (also known as the Evidence Lower Bound, ELBO) using gradient- based methods (Kingma and Welling, 2013). With Jensen's 's inequality, the variational bound is written (see the Appendix for the full derivation)
log 2 (M: Tog
Egu(scl po(orsrolse relse,
Eg 1hee- 1lhe-
KL KL(go(selhe oz) (sclhee))]]
reconstraction
regularization
Here, we denote the objective function as CMLE For each time step, the variational bound consists of two terms, i.e. the reconstruction accuracy of o, and rt, and the KL-divergence between the variational posterior and the predictive prior regularization. During model learning, the RSSM learns the latent dynamics of the environment from real trajectories by predicting future observations and rewards based on past observations and actions. During behavior learning, the RSSM generates trajectories by unrolling the state vectors using the deterministic and stochastic state modules given actions from an agent, allowing the agent learn task behavior from simulated experiences the compact latent state space.
3

---

2.2 The task-aware model loss
The RSSM is variational world model trained through MLE, which is designed to reconstruct all information in environmental observations, without taking into account the specific task of the agent. This can lead to a gap between maximizing model return and actually maximizing environment return. To understand this, consider given state-action pair (s, a) and typical actor-critic RI. agent. If the consequent state predicted by the model is off by e from the ground truth s', the state value predicted by the critic will also have error (s'' + +e) (ss)). ). As the actor's objective is to maximize state values, such errors in the predicted state values can lead to a sub-optimal policy. Therefore, if we want to truly align the model and environment for specific task, it is intuitive to evaluate model's performance using some type of value-relevant metric.
Different from MLE, Farahmand et al. (2017) proposed Value- Aware Model Learning (VAML), a loss that evaluates a model's performance by the impact of model errors on the value estimation accuracy. Given an environment transition distribution and its model approximation P. distribution over the state-action space μe and a value funetion V, the VAMI. loss is written as
CVAML (P. ((s,a)
/ p(s"|s, ((s))ss,))
d(s,a)
(3)
environment value estimate
model valuic estimate
The primary concept behind VAML is to penalize model based on the disparity between the state values of predicted states and those of the ground truth states. However, this loss function has several critical limitations when it comes to implementation on variational world models like the RSSM. In particular, VAML requires a value function that works in a pre-defined state space or the original observation space, while in DreamerV2, the critic from the agent works in the compact latent state space of the RSSM. Secondly, resolving or approximating the expectations/integral realistic RL environment, where the state space is often vast and continuous. Additionally, VAML compels the model to conserve only the value-relevant characteristies, leading to the rejection of considerable amount of semantic information contained in the environmental observations.
challenging in
To construct a task-aware model loss suitable for evaluating the RSSM in the latent state space, we explicitly distinguish prior and posterior states, and replace the environment value estimation with values from the inferred posterior states, model value estimation with values from the predicted prior states, and the expectations with an empirical summation over a given dataset resulting in parsimonious task-aware loss function, which we name V-VAM (Variational VAML)
V-VAML (T: (. = Σ V(S - (isprior)
(4)
Here, we introduce superscripts to differ posterior and prior states. Notice the states outputs of the model, and depend on the model's parameters 0 deterministically through reparameterization, i.e. sp (0ge, aee; 0) and s
prior (o<t; a.ce: 0). We leave out the inputs for simplicity.
The intuition of V-VAML loss is quite similar to the original VAML loss, yet far easier to implement given dataset. It is also evident that -V-A has some similarities with the KL regularizer in CMLE, where large distance between the posterior and prior is undesirable in both objectives. The difference is that L V-AML evaluates the impact of such distance using the value disparity. One could try somehow replacing the KL regularizer with L (-VML to achieve task-aware model learning. We take a different approach and leave that to future works.
2.3
The bi-level framework
We now introduce our proposed meta-weighting mechanism and bi-level framework to subtly fuse the RSSM and our task-aware loss function into a hierarchical optimization paradigm.
Our intention is to attain task awareness in model learning, while still maintaining an MLE foundation to preserve abundant semantic information. Inspired by advances in meta-learning (Nichol al., 2018; Shu et al., 2019; Jiang et al., 2022), we propose to assign each training sample with different

---

Algorithm 1: Task-aware Environment Modeling Pipeline with bi-level Optimization (TEMPO)
Input: Environmental trajectory world model Mo, meta weighter Na. value function V, update steps K, learning rate for model 1/, learning rate for meta weighter À while not converged do
2
for update step k = 1...... K do
Infer states with world model IT,, 11T SIT
post
prior
Mo(T);
Compute values {V (sc P V(s poir
Jtml
Normalize values (V(s, post ( TEl + Norm (V(s)ai)) ((s )) )===1);
peior IT
Compute meta weights wi.7 (Na (h,, of :66
prior
((s (s.p0)) (s prior) 1==1
Belon (TiW).T.9).
Update world model <- 0 +
80
Update meta weighter a + a ac (z:V.))
0² Jona (r:w1.T.8)
88
80 do
end
end
importance weight in a task-aware fashion. To achieve this, we introduce an additional meta weighter network N., (with parameters that outputs meta weight to for each (Oga, act))) in CMLE, based on the current states and values (Figure I(a))
Na (he, st sp (ss (sprior)
prior
(5)
With these meta weights, the model objective changes from CAL to a weighted sum of ELBO,s. The meta weights as well as the meta weighter trained to minimize task-aware model loss CV-VAMI, This adds another level of optimization on top of the default MLE model learning, forming bi-level framework with two hierarchical optimizations orthogonal objectives, which we name Task-aware Environment Modeling Pipeline with bi-level Optimization (TEMPO)
min Lupper (T; (V VAMAL (T:V Σ v(s (ogt, - V (sfrrio (s.tt (ctt,
8. L. 0" arg max Clower (T:W1:T: 0) arg max
wt (osh
(6)
Notice that l* depend on meta weights W1.T and parameter a through the gradients of Clower, Under this bi-level model learning framework, the meta weighter learns to assign importance weight to training samples regarding their impact on minimizing the task-aware model loss; and the model learns the environment dynamics through reconstruction while focusing on important training samples.
In practice, as the task behavior of the agent gradually improves, the state values naturally rise, which will likely lead to larger differences between the posterior and prior state values. If we naively optimize the meta weighter minimize Lupeer for the upper-level optimization, the meta weighter may learn to output minimal weight for all training samples to stop the values from rising. Since we are essentially interested in the relative importance of the samples, we perform normalization with moving mean and variance on the state values (V(s)ea) V(s grrirr)) before computing Lupper
To implement our TEMPO framework, we alternate between upper-level and lower-level optimization. That is, we first perform the lower-level optimization on the model parameters through single gradient update (Figure ((b)
8' = 0 + 17 V flower (T: WU.T, 0)
(7)
where 7 denotes the learning rate for the model. Then, we perform the upper-level optimization on the meta weighter parameters a, for which we can simply calculate Cupper using and take the gradient w.r.t a
5

---

(a)
Cartpole Swingup
Finger Spin
Hopper Stand
t= 3.322
p= 0.080
15.95.8 0.04.
0,787 p= 0.306
1000
1000
1000
Evaluation Return
800
800
800
6000
600
500
(b)
4000
400
400
200
200
10
10
Environmental Steps
1e5
TEMPO
Dreamerv2
D4PG (1e8) A3C (168)
Figure 2: Evaluation of TEMPO continuous control tasks from DeepMind Control Suite. (a), graphical demonstration the environments used in the evaluation, from left to right: Acrobot, Cartpole, Cheetah, Cup, Finger, Hopper. Pendulum, Quadruped, and Walker. (b), the evaluation curves of seeds and their average significance, with lines showing the mean scores, and shaded areas showing the standard deviations.
p (r; V, V,8") 00'
VaL (T: V.0")
00
da
(8)
Peerer (r; V, 0²) 82 Clower (T; WI.T. 0)
00
00 Oo
However, the problem with this direct approach is that we need to infer the states (i.e. h1.7, post prior) twice, time for each objective, for they come from different model parameters (i.e. 6 and f¹). Therefore, to cut down the computational cost, we propose to empirically approximate with 6, since they are only one update away. This results in an elegant update formula for (Figure 1(c))
a' = X V Lupper (T; V.
d (T; V,0) 8 Clower (T; W1.T. 0)
00
08 da
is evident that, with Equation 9, the upper-level objective and lower-level objective are both calculated with the same set of states, which allows us infer the states only once during epoch of bi-level optimization. See Algorithm for pseudocode of our bi-level model-learning paradigm and Figure I for graphical demonstration.
Experiments
3.1 Experimental setup
We implement our TEMPO framework on top of DreamerV2 using the official implementation of Hafner et al. (2020). Specifically, we build the meta weighter to be 5-layer dense network (MLP) with concatenated states as input (see Equation dimensions of the network are set to 400. Batch normalization (loffe and Szegedy, 2015) and ELU (Clevert et al., 2015) activation are performed after each hidden layer. sigmoid function and an additive bias are applied to the meta weights after the output layer, so that the weights center around 1, i.e. weight = 0.5 x ar (output) + 0.75. An Adam optimizer (Kingma and Ba, 2014) with learning rate of 1e-4 is used for updating the meta weighter. We leave the RSSM, the actor-critic agent, as well as all related settings from DreamerV2 untouched and fix the above configuration across all following experiments.
and scalar meta weight as output. All hidden

---

(a)
Bank Heist
Kangaroo
Hero
1.29.29 p 0.125
1.046 p= 0.258
1= 1.596
p 0.1.46
Evaluation Return
1200
w/ww
8000
14000
6000
800
10000
(b)
6000
400
2000
01
10
10
Environmental Steps
1e6
TEMPO
Dreamerv2 ION (1e7) Rainbow (167)
Figure 3: Evaluation TEMPO on discrete control tasks from Atari video games. (a), a graphical demonstration of the 6 Atari games used in the evaluation, from left to right: Bank Heist, Crazy Climber, Freeway, Hero, Kangaroo, and Pong. (b), the evaluation curves seeds and their average significance, with lines showing the mean scores, and shaded areas showing the standard deviations.
We use 9 continuous control tasks from the DeepMind Control (DMC) Suite (Tassa et al., 2018) (i.e. Acrobot Swingup, Cartpole Swingup, Cheetah Run, Cup Catch, Finger Spin, Hopper Stand, Pendulum Swingup. Quadruped Walk, and Walker Walk), and 6 discrete control tasks from Atari video games (Bellemare et al., 2013) (i.e. Bank Heist, Crazy Climber, Freeway, Hero, Kangaroo, and Pong), to evaluate the performance of TEMPO.
We pit TEMPO against the original DreamerV2 and number of state-of-the-ard model-free RL agents including D4PG Marthh-Maoo et al., 2018), A3C (Mnih et al., 2016), IQN (Dabney et al., 2018), and Rainbow (Hessel et al., 2018). For the model-free agents, Tassa et al. (2018) and Castro et al. (2018). For Dreamer V2, we use the default hyperparameters provided by Hafner al. (2020) in all experiments. More importantly, we fix these hyperparameters in the corresponding parts of our TEMPO framework (i.e. the RSSM and the actor-critic agent) for a fair comparison.
the results reported by
We tuned TEMPO on DMC Walker Walk and Atari Pong such that meta weighter converges gradually alongside the RSSM. This is to ensure the weighter doesn't 't reach early convergence Empirically, TEMPO trains stably, and both objectives reach convergence after sufficient training (Figure 4). Refer to the Appendix for further detailed settings.
Upper Obj
Cartpole Swingup
Lower
Obj
meaningless state representation.
lc4
-1.08
0.4
-1100
Following Henderson et al. (2018), we perform signif- icance testing to examine TEMPO's superiority. par- ticular, we perform the one-tailed Welch's s I-test (Welch, 1947) on TEMPO's s and Dreamer results for every evaluation during the training process, and report the average and p over all evaluations.
0.0
-.12
0
Environmental Steps
1e5
Upper objective (task aware)
Lower objective (MLE)
Figure 4: Training curve TEMPO on Cart- pole Swingup from DeepMind Control Suite. The lines show the means and shaded areas show the standard deviations of 3 seeds.
Furthermore, to verify the effectiveness and novelty of the proposed weighting mechanism, we com- pare TEMPO's meta-weighting mechanism with naively weighting the training samples using the task-aware model losses. We also perform ablation studies on dif- ferent inputs the meta weighter network and different hyprpparameters.

---

Cartpole Swingup
Hopper Stand
Evaluation Raturn
600
6000
4000
400
200
200
TEMPO
Naive Weighting
Environmental Steps
1e5
Cartpole Swingup
Hopper Stand
1000
1000
Evaluation Return
800
800
600
600
400
400
D+S+V (TEMPO)
D+S
10
10
Environmental Steps
1e5
Figure 5: Ablation studies. The lines show the mean scores and shaded areas show the standard deviations of 3 seeds. (a), comparison with naive weighting strategy. (b), Comparison of different meta weighter inputs. Capital letters denote different parts of input: deterministic states (D), stochastic states (S), and state values (V).
3.2 Continuous control
subset of the DeepMind Control Suite with tasks is used evaluate TEMPO's performance in continuous control situations. The tasks are each from different environment, as illustrated in Figure 2(a). Environmental observations are RGB images of shape 64 64 3; actions range from to 12 dimensions; each episode starts with randomized initial state and lasts for 1000 steps. We follow the protocol described by Hafner al. (2020) and set the stochastic latent variable (i.e. 8e) of DreamerV2 and TEMPO to be 32-dim continuous vector following Gaussian distribution.
We evaluate the TEMPO agent's 's environment return every 1e4 environmental steps, and compare the results with those of DreamerV2 V2 and the final performance of D4PG and A3C after 1e8 steps. The results of tasks are illustrated in Figure 2(b), where TEMPO achieved sate-of-thhe-ar performance in all tasks. Specifically, TEMPO reached or exceeded the performance of D4PG and A3C within only 1e6 environmental steps, far less than 1e8 steps. Moreover, TEMPO exceeded DreamerV2 in terms of asymptotic performance, training stability, and convergence speed, which demonstrates the advantage of TEMPO's task-aware model learning paradigm under continuous task settings. Refer to the Appendix for full results on all tasks.
3.3 Discrete control
We then evaluate TEMPO's performance on discrete control tasks with 6 Atari video games, as illustrated Figure 3(a). The actions range from to 18 dimensions. We render the environmental observations gray-scale images of shape 64 64 x 1, and set the stochastic latent variable (i.c. st) to be a 32-column discrete matrix with each column being a 32-dim one-hot vector from a categorical distribution, following Hafner et al. (2020).
Due to limited computational resources, we train TEMPO and DreamerV2 for 1c7 environmental steps. We evaluate the TEMPO agent's 's environment return every 1c5 steps, and compare the results with those of DreamerV2 and the final performance of IQN and Rainbow after 1e7 steps. The results of 3 games are shown in Figure 3(b), where TEMPO again achieved top performance and exceeded the original DreamerV2. This demonstrates TEMPO's superiority and robustness under both continuous and discrete task settings. Refer to the Appendix for full results on all 6 games.

---

Cartpole Swingup
800
Evaluation Return
600
200
10
Environmental Steps
1e5
DreamerV2
TEMPO (hidden_dimg clim 200)
TEMPO (activation ReLU)
TEMPO
TEMPO (layers 3)
TEMPO (norm LLaeeNorm)
Figure 6: Comparison of different hyperparameters of the meta weighter. The lines show the mean scores and shaded areas show the standard deviations of 3 seeds.
3.4 Comparison with naive weighting strategy
In TEMPO, the meta weighter, deep neural network, generates the meta weights based state information. Given that the meta weighter is trained to minimize our task-aware model loss, one may naturally wonder:
viable to naively use the losses themselves sample weights, assigning greater weights to samples with larger value disparities? To compare this straightforward weighting approach with our TEMPO framework, we initially normalize the task-aware model losses of each time step in training trajectory using moving mean and variance. Subsequently, we add I to these normalized losses obtain the final naive sample weights, roughly centering the weights around 1. As illustrated in Figure 5(a), the naive weighting strategy exhibits strong instability during training, and fails gain comparable results as TEMPO.
3.5 Comparison of different meta weighter inputs
In our initial design, the meta weighter generates sample weights by considering three key components of information: deterministic state h, stochastic states sp and prior and state values (s and ((stter) (see Equation 5). To justifyy our design, conduct ablation study on the inputs to assess the impact of these inputs on agent's performance. We sequentially remove the state values and deterministic state from the input, and compare the outcomes with original TEMPO configuration. As illustrated in Figure 5(b), removing the two parts of input causes the agent to converge slower, and introduces additional instability and variance in the agent's performance. Notably, removing the deterministic states severely affects the outcome performance. This demonstrates the significance of the information encapsulated within these deterministic states.
3.6 Comparison of different hy perppaaameters
To examine the robustness TEMPO's meta-weighting mechanism, we perform simple ablation studies on number of important hyperparameters of the meta weighter, including hidden dimension, the number of dense layers, the choice of nonlinear activation function, and the normalization method. The results, depicted in Figure 6, clearly demonstrate TEMPO's robustness, especially in face of reduced hidden dimensions and altered activation functions. While modification to normalization method and decrease in the number of layers does affect performance, TEMPO can consistently maintain an advantage over DreamerV2 during the initial stages of training, ensuring quicker convergence. Overall, TEMPO proves to be highly adaptable various hyperparameters.

---

4
Related work
Model learning in partially observable MBRL problems has drawn large interest in related works. A straightforward way for model learning is to learn the dynamics of an environment by fitting the observations from before and after environment transitions, either deterministically or stochastically (Oh et al., 2015; Chiappa al., 2017; Kaiser et al., 2019). We call this type of model observation space models. Simulated trajectories can be obtained by unrolling the model in the original observation space, which facilitates behavior learning in Dyna fashion (Sutton, 1991). For example, SimPL.e (Kaiser et al., 2019) proposed a video prediction network for pixel-level prediction on environmental observations, where a PPO (Schulman et al., 2017) agent is trained using the predictions achieve behavior learning for Atari games.
However, modeling high-dimensional image observations is unavoidably computationally intensive. Latent state-space models overcome this limitation by constructing a compact latent space to char- acterize the environmental states behind each observation, which minimizes the memory footprint during model unrolling (Watter et al., 2015; Wahlström et al., 2015; Buesing et al., 2018; Gelada et al., 2019; Hafner al., 2019b,a, 2020; Ozair al., 2021). Based the RSSM, Dreamer (Hafner et al., 2019a) and V2 (Hafner et al., 2020) enable behavior learning by constructing an actor-critic agent that maximizes state values by propagating their analytic gradients back through the dynamics learned by the model. Ozair et al. (2021) proposed a stochastic state-space model based on VQ- (Van Den Oord al., 2017), which is then combined with Monte Carlo tree search (MCTS (Coulom, 2007)) for task behavior.
Observation-spent spacin models and state-space models both be categorized MLE-based models that learn environment dynamics through reconstruction. With these models, model errors are likely to compound during unrolling. Value-equivalent models intend to minimize such errors and improve task awareness by focusing on end-to-ent value prediction (Tamar et al., 2016; Silver et al., 2017; Farquhar al., 2017; Farahmand, 2018; Hubert et al., 2021). Specifically, methods like the Predictron (Silver et al., 2017), Value Prediction Network (VPN (Oh et al., 2017)), MuZero (Schrittwieser et al., 2020), and Stochastic MuZero (Antonoglou et al., 2021) aims to form state representation that can help a parametric value function make accurate value predictions. Other methods like VAML (Farahmand et al., 2017) and VaGraM (Voelcker et al., 2022) design task/value-aware model losses that minimize value prediction disparities from a given value function.
Although we base TEMPO on the RSSM and VAML in this work, TEMPO is orthogonal to the design of MLE- world models (lower level) and task-aware model loss (upper level). On one hand, one can find suitable task-aware model loss to evaluate a MLE-based model (e.g. VAML for observation-space models and our V-VAMI. for state-space models), and build a bi-level optimization to train the model in TEMPO fashion. On the other hand, one can replace our V-VAML loss with another task-aware model loss (e.g. VaGraM) by making modifications similar to V-VAM..
5 Limitation and discussion
We present TEMPO, a task-aware framework for learning world models. The framework is driven by a meta-weighting mechanism and novel task-aware loss function under a bi-level optimization. TEMPO achieves state-of-the-art performance variety of continuous and discrete control tasks, while demonstrating better asymptotic performance, training stability, and convergence speed. Abla- tion studies are performed to justify method. The main limitation of our work is that TEMPO is more computationally demanding than previous model learning methods, as it involves two loops of optimization. In our experiments, observed that TEMPO trains around 40% slower and requires 80% more RAM than V2.. Nonetheless, TEMPO contributes significant insights and opens up an exciting new avenue for environment modeling- We aspire that TEMPO will serve catalyst for novel ideas and approaches for masel- RL in the future. A sample code of TEMPO is available at https: //tp:/tgtttttttt/ com/deng-ai -aann--a--abbTEee lab/TEMPO.
Acknowledgements
This work was supported in part by the National Natural Science Foundation of China under Grant 62325101 and Grant 10031001.
10

---

References
loannis Antonoglou, Julian Schrittwieser, Sherjil Ozair, Thomas Hubert, and David Silver. Planning in stochastic environments with learned model. In International Conference Learning Representations, 2021.
Gabriel Barth-Maron, Matthew W Hoffman, David Budden, Will Dabney, Dan Horgan, Dhruva Tb, Alistair Muldal, Nicolas Heess, and Timothy Lillicrap. Distributed distributional deterministic policy gradients. arXiv preprint 18x.v:/804.08611 2018.
M. G. Bellemare, Y. Naddaf, J. Veness, and M. Bowling. The arcade learning environment: evaluation platform for general agents, Journal of Artificial Intelligence Research, 47:253-279, jun 2013.
Lars Buesing, Theophane Weber, Sébastien Racaniere, SM Eslami, Danilo Rezende, David P Reichert, Fabio Viola, Frederic Besse, Karol Gregor, Demis Hassabis, et al. Learning and querying fast generative models for reinforcement learning. arXiv preprint 1oXiv:/802.030000 2018.
Pablo Samuel Castro, Subhodeep Moitra, Carles Gelada, Saurabh Kumar, and Marc G Belle- mare. Dopamine: A research framework for deep reinforcement learning. arXiv preprint
arXiv:1812.00110, 2018.
Silvia Chiappa, Sébastien Racaniere, Daan Wierstra, and Shakir Mohamed, Recurrent environment simulators. arXiv preprint 3xXiv:/704.02244 2017.
Djork-Aré Clevert, Thomas Unterthiner, and Sepp Hochreiter. Fast and accurate deep network learning by exponential linear units (elus). arXiv preprint arXiv:1511.07289, 2015.
Rémi Coulom. Efficient selectivity and backup operators in monte-carlo tree search, In Computers and Games: 5th International Conference, CG 2006, Turin, Italy, May 29-31, 2006. Revised Papers 5, pages 72-83. Springer, 2007.
Will Dabney, Georg Ostrovski, David Silver, and Rémi Munos. Implicit quantile networks for distributional reinforcement learning. In International conference on machine learning, pages 1096-1105. PMLR, 2018.
Amir-massoud Farahmand. Iterative value-aware model learning. Advances in Neural Information Processing Systems, 31, 2018.
Amir-massout Farahmand, Andre Barreto, and Daniel Nikovski. Value-aware loss function for model-based reinforcement learning. In Artificial Intelligence and Statistics, pages 1486-1494. PMLR, 2017.
Gregory Farquhar, Tim Rocktäschel, Maximilian Igl, and Shimon Whiteson. Treeqn and atreec: Differ- entiable tree-structured models for deep reinforcement learning, arXiv preprint arXiv:1710.114:: 2017.
Carles Gelada, Saurabh Kumar, Jacob Buckman, Ofir Nachum, and Marc G Bellemare. Deepmdp: Learning continuous latent space models for representation learning. In International Conference on Machine Learning, pages 1177--19, PMLR, 2019.
Christopher Grimm, André Barreto, Satinder Singh, and David Silver. The value equivalence principle for model-based reinforcement learning. Advances in Neural Information Processing Systems, 33: 2020.
Christopher Grimm, André Barreto, Greg Farquhar, David Silver, and Satinder Singh. Proper value equivalence. Advances in Neural Information Processing Systems, 34:7773-7786. 2021.
David Ha and Jürgen Schmidhuber. Recurrent world models facilitate policy evolution. Advances in neural information processing systems, 31, 2018
Danijar Hafner, Timothy Lillicrap, Jimmy Ba, and Mohammad Norouzi. Dream to control: Learning behaviors by latent imagination. arXiv preprint arXiv:1912.01603 2019a.

---

Danijar Hafner, Timothy Lillicrap, lan Fischer, Ruben Villegas, David Ha, Honglak Lee, and James Davidson. Learning latent dynamics for planning from pixels. In International conference on machine learning, pages 2555-2565. PMLR, 2019b.
Danijar Hafner, Timothy Lillicrap. Mohammad Norouzi, and Jimmy Ba. Mastering atari with discrete world models. arXiv preprint crXiv:2010.02/993 2020.
Peter Henderson, Riashat Islam, Philip Bachman, Joelle Pineau, Doina Precup, and David Meger. Deep reinforcement learning that matters. In Proceedings of the AAAI conference on artificial intelligence, volume 32, 2018.
Matteo Hessel, Joseph Modayil, Hado Van Hasselt, Tom Schaul, Georg Ostrovski, Will Dabney, Dan Horgan, Bilal Piot, Mohammad Azar, and David Silver. Rainbow: Combining improvements in deep reinforcement learning. In Proceedings of the AAAI conference on artificial intelligence. volume 32, 2018.
Thomas Hubert, Julian Schrittwieser, Ioannis Antonoglou, Mohammadamin Barekatain, Simon Schmitt, and David Silver. Learning and planning in complex action spaces. In International Conference on Machine Learning, pages 4476-4486. PMLR, 2021.
Sergey Ioffe and Christian Szegedy. Batch normalization: Accelerating deep network training by reducing internal covariate shift. In International conference on machine learning, pages 448-456. pmlr, 2015.
Michael Janner, Justin Fu, Marvin Zhang, and Sergey Levine. When to trust your model: Model-based policy optimization. Advances in neural information processing systems, 32, 2019.
Xingyu Jiang, Hongkun Dou, Chengwei Fu, Bingquan Dai, Tianrun Xu, and Yue Deng. Boosting supervised dehazing methods via bi-level patch reweighting. In Computer Vision-ECCY 2022: 17th European Conference, Tel Aviv, Israel, October 23-27, 2022, Proceedings, Part XVIII, pages 57-73. Springer, 2022.
Lukasz Kaiser, Mohammad Babaeizadeh, Piotr Milos, Blazej Osinski, Roy H Campbell, Konrad Czechowski, Dumitru Erhan, Chelsea Finn, Piotr Kozakowski. Sergey Levine, et al. Model-based reinforcement learning for atari. arXiv preprint 1:X3xx:993300077 2019.
Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint 11X.v:1412.6980, 2014.
Diederik P Kingma and Max Welling. Auto-encoding variational bayes. 1rXiv:1312.6114, 2013.
arXiv preprint
Nathan Lambert, Brandon Amos, Omry Yadan, and Roberto Calandra. Objective mismatch in model-based reinforcement learning. arXiv preprint arXiv:2002.04522: 2020.
Volodymyr Mnih, Adria Puigdomenech Badia, Mehdi Mirza, Alex Graves, Timothy Lillicrap, Tim Harley, David Silver, and Koray Kavukcuoglu. Asynchronous methods for deep reinforcement learning. In International conference on machine learning, pages 1928-1937. PMLR, 2016.
Thomas M Moerland, Joost Brockens, Aske Plaat, Catholijn M Jonker, et al. Model-based rein- forcement learning: A survey. Foundations and Trends@ in Machine Learning, 16(1):1-118. 2023.
Alex Nichol, Joshua Achiam, and John Schulman. On first-order meta-learning algorithms. arXiv preprint arXiv:1803.022 029999, 2018
Evgenii Nikishin, Romina Abachi, Rishabh Agarwal, and Pierre-Luc Bacon. Control-oriented model- based reinforcement learning with implicit differentiation. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 36, pages 7886-7894, 2022.
Junhyuk Oh, Xiaoxiao Guo, Honglak Lee, Richard L. Lewis, and Satinder Singh. Action-conditional video prediction using deep networks in atari games. Advances neural information processing systems, 28, 2015.
12

---

Junhyuk Oh, Satinder Singh, and Honglak Lee. Value prediction network. Advances in neural information processing systems, 30, 2017.
Sherjil Ozair, Yazhe Li, Ali Razavi, Ioannis Antonoglou, Aaron Van Den Oord, and Oriol Vinyals. Vector quantized models for planning. In International Conference on Machine Learning, pages 8302 8313. PMLR, 2021.
Julian Schrittwieser, Ioannis Antonoglou, Thomas Hubert, Karen Simonyan, Laurent Sifre, Simon Schmitt, Arthur Guez, Edward Lockhart, Demis Hassabis, Thore Graepel, et al. Mastering atari, go, chess and shogi by planning with a learned model. Nature, 588(7839):604-660 2020.
John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. arXiv preprint arXiv:1707.06347, 2017.
Jun Shu, Qi Xie, Lixuan Yi, Qian Zhao, Sanping Zhou, Zongben Xu, and Deyu Meng. Meta- weight-net: Learning an explicit mapping for sample weighting. Advances in neural information processing 8systems, 32, 2019.
David Silver, Hado Hasselt, Matteo Hessel, Tom Schaul, Arthur Guez, Tim Harley, Gabriel Dulac- Arool, David Reichert, Neil Rabinowitz, Andre Barreto, et al. The predictron: End-to-end learning and planning. International Conference on Muchine Learning, pages 3191-3199. PMLR, 2017.
Kihyuk Sohn, Honglak Lee, and Xinchen Yan. Learning structured output representation using deep conditional generative models. Advances in neural information processing systems, 28, 2015.
Richard S Sutton. Dyna, an integrated architecture for learning, planning, and reacting. ACM Sigart Bulletin, 2(4):160-163, 1991.
Richard Sutton and Andrew G Barto. Reinforcement learning: An introduction. MIT press, 2018.
Aviv Tamar, Yi Wu, Garrett Thomas, Sergey Levine, and Pieter Abbeel. Value iteration networks. Advances in neural information processing systems, 29, 2016.
Yuval Tassa, Yotam Doron, Alistair Muldal, Tom Erez, Yazhe Li, Diego de Las Casas, David Budden, Abbas Abdolmaleki, Josh Merel, Andrew Lefrancq, et al. Deepmind control suite. arXiv preprint arXiv:1801.00690 2018
Aaron Van Den Oord, Oriol Vinyals, al. Neural discrete representation learning. Advances in neural information processing systems, 30, 2017.
Claas Voelcker, Victor Liao, Animesh Garg, and Amir massoud Farahmand. Value gradient weighted model-based reinforcement learning. arXiv preprint vrXiv:2204.01464 2022.
Niklas Wahlström, Thomas Schön, and Mare Peter Deisenroth. From pixels to torques: Policy learning with deep dynamical models. arXiv preprint arXiv: 2015.
Manuel Watter, Jost Springenberg, Joschka Boedecker, and Martin Riedmiller. Embed to control: A locally linear latent dynamics model for control from raw images. Advances in neural information processing systems, 28, 2015.
Beraard L Welch. The generalization of 'student's' problem when several different population parobett''prooblm varlances involved. Biometrika, 34(1-2):28 35, 1947.
Amy Zhang, Rowan McAllister, Roberto Calandra, Yarin Gal, and Sergey Levine. Learning invariant representations for reinforcement learning without reconstruction. arXiv preprint 101.22006..0742,. 2020.
13

---

A Further discussion
A.1 Task selection
Our primary concern when choosing tasks is mainly the computational cost. With limited computation resources, we couldn't afford run the extensive 5e6 steps for DMC and 2e8 steps for Atari as the Dreamer papers did. Instead, we only had a 1e6-step budget for DMC and a 1e7-step budget for Atari. Consequently, the case of DMC, we focus on choosing tasks from different environments on which Dreamer can roughly reach convergence in under 1c6 steps. Similarly, for Atari, we focus on tasks that DreamerV2 could reach convergence within our 1c7 budget. We are certainly enthusiastic about exploring how TEMPO performs more challenging tasks in the future.
A.2 Network settings
The architectural configurations of our meta weighter, such as the number of layers and hidden dimensions, were informally aligned with the network modules employed in DreamerV2. Specifically, we mirrored the architectural choices found in modules like the reward module of RSSM, the actor, and the critic, all of which adopt similar MLP Multi-Layer Perceptron) structure. Given that these modules can handle state representation well, was logical choice build similar architecture for our meta weighter as well. Perhaps the most important thing we discovered about architecture design is that batch normalization in the meta weighter greatly boosts the final performance.
For the mapping the network outputs to meta weights, we toyed with weight = (outputt) + 0.5 and weight 0.5 a (output) + 0.75 initial experiments and didn't observe a noticeable difference, so we stuck with the latter in the hope of stable training.
A.3 Tuning strategy
One of the key pursuits when built TEMPO was to design paradigm that could outperform DreamerV2 with little tuning demand. Since TEMPO operates as an additive bonus algorithm, essentially fine-tuning the training sample weights within DreamerV2, it is quite easy for TEMPO to at least meet DreamerV2's performance, resulting in low tuning demand. Specifically, we tuned TEMPO on DMC Walker Walk and Atari Pong such that meta weighter converges gradually alongside the RSSM. This is to ensure the weighter doesn't 't reach early convergence on meaningless state representation. We have no doubt that superior architecture and hyperparameters exist for TEMPO, and we remain committed to further exploration and refinement.
A.4 Performance analysis
Based findings, TEMPO demonstrates notable advantage over DreamerV2 in specific tasks, such as Finger Spin, while in others, like Walker Walk, the two methods exhibit similar performance levels. We hypothesize that the distance between naive state representation from MLE and the task-aware state representation can vary on different tasks. When the distance is large, meaning that the features that can effectively lower reconstruction error don't 't coincide well with the task-specific features, DreamerV2 can stuck in learning fask-irreleve TEMPO's bi-level learning quickly captures the rask-relevant features, and lets the agent performance snowball. In other tasks, where the distance is not sigmificant, TEMPO can have less ad advantage.
features like a local minimum, while
We also observe that TEMPO shows larger variances in some tasks than others. We suspect the reason that, in these tasks, the meta weighter learns too "eagerly" relative to the RSSM, leading to less meaningful meta weights, potentially disturbing the learning of RSSM. Lowering the learning rate for the weighter may be helpful in reducing the variance in these tasks.
A.5 Future directions
World models, especially MLE-based models, encode a broad spectrum of environmental dynamics. They carry the potential of multi-tasking with a single model, presenting a possible way to multi- tasking agents, even general AI. With the hierarchical environment modeling paradigm that TEMPO offers, we aspire for it to serve as a foundational framework for enhancing task awareness across multiple specific tasks within future world models.
14

---

B Derivations
The variational bound, i.. Evidence Lower Bound (ELBO), for the Recurrent State-space Model
(RSSM) with a variational posterior (( = I ((s, Mh, or), written as
log p
/ II
log
p(pi (he) plo,, rele,relse.hy ha) dsi.7
(=1
/ III ((sslhe o.) p(sellhe) relss he) dsir
log
g(s: (1: 0r)
p(sellee) plo, rells, he) g(3://e 01)
(10)
= log (ar)
T
2 (((0nrlo11
log rils,, h1) + (og((s)))h)) loga(se
T
Egts [plon, rist, h.]] Egise 1l6, a a [ 0r) [|p(s|/hi]]]
t=1
15

---

C Hyperparameters
Module
Name
Value
World model
Batch size
16
Trajectory length
50
KL balancing
0.8
for DMC
Deterministic state dimension
200
Stochastic state dimension
32-dim continuous
Learning rate
3e-4
KL scale
1.0
for Atari
Deterministic state dimension
600
Stochastic state dimension
32-dim discrete (32 classes/dim)
Learning rate
2e-4
KL scale
0.1
Agent
Imagination horizon
15
Discount factor
0.99
Discount factor for X-target
0.95
Target critic update interval
100
for DMC
Actor learning rate
8e-5
Critic learning rate
8e-5
for Atari
Actor learning rate
4e-5
Critic learning rate
le-4
Meta weighter
Number of dense layers
Hidden dimension
400
Activation function
ELU
Normalization
Batch
Learning rate
le-
for DMC
Input dimension
200 + 2 32 + x
for Atari
Input dimension
600 + x 32 x 32 + x 2
Common
Optimizer
Adam
Gradient clipping
100
Adam epsilon
lc-5
Weight decay
1e-6
Table 1: Main hyperparameters of TEMPO. DMC stands for DeepMind Control tasks. We the default setting of Dreamerv2 in World model, Agent, and Common. All experiments were run on single Nvidia RTX 3090 GPU with Python 3.7 and Tensorflow 2.6. Refer our sample code for full implementation details.
16

---

D Full results on continuous control tasks
Acrobot Swingup
Cartpole Swingup
Cheetah Run
t= 0.940 p 0.371
1= 3.322
p 0.080
t 0.567
p 0.371
1000
1800
800
800
400
600
6001
400
400
200
100
0
10
0
2
6
10
Cup Catch
Finger Spin
Hopper Stand
t= 0.878
p= 0.410
t= 15.985
p 0.048
0.787 p 0.305
1000
1000
1000
Evaluation Return
800
800
800
600
600
600
100
400
400
200
200
200
10
Pendulum Swingup
Quadruped Walk
Walker Walk
1.004 p.
D.51g p
p 0.327
1000
800
800
600
600
600
400
400
4000
200
200
200
0
Environmental Steps
1e5
TEMPO
DreamerV2 -DPPG (1e8)
A3C (1e8)
Figure 7: Evaluation of TEMPO on continuous control tasks from DeepMind Control Suite. The lines show mean scores and the shaded areas show the standard deviation across 3 random seeds.
17

---

E Full results on discrete control tasks
Bank Heist
Crazy Climber
Freeway
3.291 0:125
1 868 p
p
1200
/ww
12
800
400
Evaluation Return
10
Hero
Kangaroo
Pong
1x3
I 1.598
148
..048 0.258
t 0.173 p 0.454
14
6000
10.
4000
2-
10
0
10
Environmental Steps
1e6
TEMPO DreamerV2 IQN (1e7)
Rainbow (1e7)
Figure 8: Evaluation of TEMPO discrete control tasks from Atari video games. The lines show
mean scores and the shaded areas show the standard deviation across random seeds.
18

---

Training curves
Upper objective (task aware)
Lower objective (MLE)
1.0
-10800
0.5
-10000
-11000
Cartpole Swingup
0.2
-11200
0.0
-11300
(a)
Value
Unweighted ELBO
-11294
1.00
-11298
100
50
-11.002
-11306
0
10
Environmental Steps
1e5
Upper objective (task aware)
Lower objective (MLE)
-3600
0.5
Bank Heist
40000
(b)
Value
Unweighted ELBO
40
-3800
30
-3850
40000
Environmental Stops
106
TEMPO DreamerV2
Figure 9: Training curves of TEMPO on (a) Cartpole Swingup from DeepMind Control Suite and (b) Bank Heist from Atari, including curves of value estimation, upper objective, lower objective, and unweighted ELBO. The lines show means and the shaded areas show the standard deviation across random seeds.
19

---

G
Sample weights
Set
Type
Weights (W1:50)
Pearson (r)
1
meta
[1.18
1.20
1.20
1.16
1.20
1.04
0.99
0.95 ...]
-0.20
naive
[1.57
0.27
0.35
1.16
0.52
1.07
0.43
0.78
...]
2
meta
[1.17
1.06
1.08
1.13
0.97
0.93
1.03
1.16
...]
-0.09
naive
[0.29
1.17
0.68
0.43
0.35
1.67
1.70
0.70
...]
meta
[1.07
1.09
0.88
0.80
1.10
0.91
0.94
0.93
0.29
naive
[2.69
4.74
1.42
1.14
2.39
0.49
0.44
0.39
...]
meta
[1.09
1.13
1.01
1.17
1.14
1.09
0.92
0.96
...]
0.02
naive
[0.27
0.28
0.79
0.27
0.34
0.28
0.57
0.32
Table 2: 4 sets of meta weights and corresponding naive weights after 5e5 steps of training in DMC Cartpole Swingup. The Pearson correlation coefficients show that there is apparent linear relation between the two types of weights.
20


Reviews:
[{"review_id": "T4noFwy0AW", "reviewer": "Reviewer_9zSq", "summary": "The authors propose to combine the benefits of value-based model-based methods (like MuZero) with reconstruction-based model-based methods (like Dreamer). They do so through TEMPO, which is a bi-level model-learning algorithm that includes a meta weighter network that weights each training sample with the goal of minimizing a task-aware loss. They present experimental results to support the claims of state-of-the-art performance.", "questions": "1. Could you not just use a smaller representational capacity since the task-awareness should bias your network toward preserving the most task-relevant information? Or am I missing something?\n\n2. Why do you think TEMPO exhibits more training variance on only some environments (e.g., Figure 2)? And why do you think that some environments yield similar performance between Dreamerv2 and TEMP (e.g., Hopper Stand) while others have a large discrepancy (e.g., Finger Spin)? \n\n3. Given that TEMPO is more computationally demanding, are there any optimization tricks that you believe could help? \n\n4. Is the reason for performing this optimization in a bi-level manner to not force the model to trade-off between preserving value-relevant information and performing reconstruction? The trade-off between reconstructing the different components of the MDP seems to be a major bottleneck in Dreamer, as one must preselect a network capacity that can preserve all components in the hopes that one will capture the value-relevant ones. ", "limitations": "The paper includes a nice discussion section that includes the main limitation of the work. Although it would be nice to see a more detailed discussion of limitations, I recognize that authors are provided limited space to include all contributions, so I treat the inclusion of any limitation as a bonus toward the paper.", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The proposed idea of combining the best of both worlds with respect to model-based methods is interesting and timely. The idea of using the upper level to weight samples is also of interest.\n\n2. I like that the authors are clearly quite familiar with the related literature. I appreciate the effort and detail provided to summarize the most relevant work in a way that could help even a newcomer to the field understand the current state of the art. \n\n3. In general, I think this paper is well-written and clear. Similar to my comment in strength 2, I think the authors do a nice job of explaining their method in an accessible way. \n\n4. I appreciate that the authors included a concrete discussion section that explicitly mentions the weaknesses of TEMPO. \n\n5. The authors show promising empirical results: the underlying Dreamer-v2 model for TEMPO does not need any additional hyperparameter tuning to enable TEMPO to achieve performance that is at least as good as the vanilla Dreamer-v2 algorithm. ", "weaknesses": "1. I encourage the authors to make more clear what is their contribution and what is contributed by previous work in Section 2. It is well-written in general, but including things like the details about deriving the variational bound for ELBO may be unnecessary? \n\n2. I would like to see proper significance testing to support the claims of superiority in the paper [Henderson et al., 2018]. I also would like to see a mention of how the presented seeds were chosen for the environments. I would also like to see how the hyperparameters for TEMPO were chosen. I see that the RSSM + actor-critic were untouched, but what about the other architectural and optimization choices (e.g., Adam learning rate)? How were they selected?\n\n[Henderson et al., 2018] Henderson, Peter, et al. \"Deep reinforcement learning that matters.\" Proceedings of the AAAI conference on artificial intelligence. Vol. 32. No. 1. 2018.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Dear Reviewer 9zSq,\n\nWe sincerely appreciate your positive comments and feedback, as well as the suggestions for improvement. In response to your questions, we have prepared the answers below:\n\n---\n**Q1: I encourage the authors to make more clear what is their contribution and what is contributed by previous work in Section 2. It is well-written in general, but including things like the details about deriving the variational bound for ELBO may be unnecessary?**\n\n**A1:** This is a valid concern. The RSSM world model in Section 2.1 is the work of Dreamer. I suppose this can lead to some confusion about our contribution. We will try to highlight our contribution more in the revised version. On the other hand, you could say it is unnecessary to include the derivation of ELBO for RSSM, but we just want to make the paper as self-contained as possible to reduce readers' burden of referring to other works in order to understand ours.\n\n---\n**Q2: I would like to see proper significance testing to support the claims of superiority in the paper [Henderson et al., 2018].**\n\n**A2:** Following [Henderson et al., 2018], we further conducted a number of significance testing experiments regarding different hyperparameters for the TEMPO meta weighter, including hidden dimension size, number of layers, activation function, and normalization. Other settings remain the same to maintain a fair comparison with the original DreamerV2. The curves are presented in Figure 1 of the rebuttal pdf.\n\nAlthough showing sensitivity to the number of layers and type of normalization, TEMPO is quite robust overall.  In fact, one of the key pursuits when we built TEMPO is to design a paradigm that can outperform DreamerV2 with little tuning demand. With TEMPO being essentially an additive bonus algorithm that weights DreamerV2's training sample, it is quite easy for TEMPO to at least meet DreamerV2's performance, leading to low tuning demand.\n\n---\n**Q3: I also would like to see a mention of how the presented seeds were chosen for the environments.**\n\n**A3:** The 3 random seeds were simply fixed to be 0, 1, and 2 in all of our experiments.\n\n---\n**Q4: I would also like to see how the hyperparameters for TEMPO were chosen. I see that the RSSM + actor-critic were untouched, but what about the other architectural and optimization choices (e.g., Adam learning rate)? How were they selected?**\n\n**A4:** As we mentioned in A2, we built TEMPO to have low tuning demand. Architectural settings like the number of layers and hidden dimensions were quite casually set to be the same with network modules from DreamerV2 like the reward module of RSSM, actor, and critic (these modules all have similar MLP architecture). Since the network architecture in these modules can handle state representation well,  we can naturally use it for our meta weighter as well. Perhaps the most important thing we discovered about architecture design is that batch normalization in the meta weighter greatly boosts the final performance. \n\nAs for the learning rate, we tuned it on DMC Walker Walk and Atari Pong such that meta weighter converges gradually alongside the RSSM. This is to ensure the weighter doesn't reach early convergence on a meaningless state representation. We have no doubt that better architecture and hyperparameters exist for TEMPO.\n\n---\n**Q5: Could you not just use a smaller representational capacity since the task-awareness should bias your network toward preserving the most task-relevant information? Or am I missing something?**\n\n**A5:** Intuitively, yes. We wanted TEMPO as a performance boost to DreamerV2, so we try not to (or rather didn't have the time) poke around with the DreamerV2 part of our algorithm. After all, DreamerV2 has a lot of settings and hyperparameters. We were worried that reducing the model size would result in the need the re-tune other hyperparameters (e.g. learning rate) as well.\n\n---\n**Q6: Why do you think TEMPO exhibits more training variance on only some environments (e.g., Figure 2)? And why do you think that some environments yield similar performance between DreamerV2 and TEMP (e.g., Hopper Stand) while others have a large discrepancy (e.g., Finger Spin)?**\n\n**A6:** This is a great question. We suspect the reason that TEMPO shows larger variances in some tasks is because, in these envs, the meta weighter learns too \"eagerly\" relative to the RSSM, leading to less meaningful weights, potentially disturbing the learning of RSSM. Lowering the learning rate for the weighter may be helpful in reducing the variance in these tasks. \n\nWe believe that in some tasks, like Finger Spin, there is a significant distance between naive state representation from MLE and the task-aware state representation, and DreamerV2 can stuck in learning task-irrelevant features like a local minimum, while TEMPO's bi-level learning quickly captures the task-relevant features, and lets the agent performance snowball. In other tasks, where the distance is not significant, TEMPO can have less advantage.\n\n---\n**Q7: Given that TEMPO is more computationally demanding, are there any optimization tricks that you believe could help?**\n\n**A7:** Since the initial submission, we have optimized our code, resulting in a much faster implementation of TEMPO. The new code, while still consuming around 80\\% more RAM than DreamerV2, runs only less than 40\\% slower than DreamerV2. We will update this number in the revised version.\n\n---\n**Q8: Is the reason for performing this optimization in a bi-level manner to not force the model to trade-off between preserving value-relevant information and performing reconstruction?**\n\n**A8:** Yes, this trade-off is exactly our motivation. An MLE-based model like RSSM dedicates all of its network capacity to reconstructing all information equally. With TEMPO, We want to hold out some of the network capacity and dedicate them to focusing on reconstructing task-relevant information. So, yes, it is an intended trade-off."}, {"review_id": "2BMsF8Xids", "reviewer": "Reviewer_kk1X", "summary": "This paper proposes a novel model-based algorithm termed Task-aware Environment Modeling Pipeline with Bi-level Optimization (TEMPO). To achieve a proper trade-off between MLE and Value-Equivalence models, TEMPO utilizes a bi-level optimization framework that trains an upper-level weight network generating non-uniform sample weights concerning the value equivalence and a low-level RSSM learning the state representations. The performance of TEMPO is examined in benchmarks involving continuous control tasks and video games.", "questions": "N/A", "limitations": "N/A", "rating": 5, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "- The paper is well-organized and the proposed method is easy to understand.\n\n- The experiments are conducted in more than one particular benchmark.", "weaknesses": "One concern is that the proposed method is not well-motivated, and there are some factual errors in the texts about motivation. As claimed in lines 38-42 and 115-11, the MLE-learned representation overlooks the task-specific information. Actually, the learned representations of Dreamerv2 not only rely on the reconstruction loss but the reward prediction and transition prediction that indeed provides task-specific information. Nevertheless, the claim made in lines 130-131 is not necessarily a limitation of VAML, since we can simply use some existing representation learning technique to convert the obs in POMDP to a low-dimensional state, in which case VAML loss is still applicable. \n\nThe experiments with only 3 random seeds are not convincing to me. Only 3/9 results of DMC and 3/6 results of Atari are demonstrated in the main paper, which is not enough. Even though the results in the remaining environments are not promising, the main paper should demonstrate most of them.\n\nTaking the significant computational demand of TEMPO (lines 307-308) into account, I think the demonstrated performance improvement is marginal.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Dear Reviewer kk1X,\n\nThank you for your detailed and constructive feedback. Here is our response to your concerns:\n\n---\n**Q1: The learned representations of DreamerV2 not only rely on the reconstruction loss but the reward prediction and transition prediction that indeed provides task-specific information.**\n\n**A1:** In common RL envs, state transition prediction only provides general information about the fundamental dynamics and physics of those envs. Reward prediction does offer task-specific information. However, the action of agents, either actor-critic agents or value-based agents, are more directly influenced by the value function, which RSSM is unaware of. I suppose the phrase \"completely overlooking\" is a bit too strong. We will modify this in the revised version.\n\n---\n**Q2: The claim made in lines 130-131 is not necessarily a limitation of VAML, since we can simply use some existing representation learning technique to convert the obs in POMDP to a low-dimensional state, in which case VAML loss is still applicable.**\n\n**A2:** Yes, but doing so adds additional difficulty. Therefore, not being applicable to common situations right out of the box is a valid limitation.\n\n---\n**Q3: The experiments with only 3 random seeds are not convincing to me.**\n\n**A3:** We are well aware that using 3 random seeds is not ideal when a lot of other RL works use 5. The reason for this is simply limited time and computation resources. However, TEMPO does outperform DreamerV2 by a great margin in our experiments, which is enough to effectively support our method.\n\n---\n**Q4: Only 3/9 results of DMC and 3/6 results of Atari are demonstrated in the main paper, which is not enough. Even though the results in the remaining environments are not promising, the main paper should demonstrate most of them.**\n\n**A4:** In fact, TEMPO outperforms DreamerV2 in most of our experiments and meets DreamerV2 in the rest, despite the margin being relatively small in a few tasks. So I'm not sure that \"the remaining environments are not promising\" does justice to our work. The reason not to demonstrate all curves is simply the 9-page space limit.\n\n---\n**Q5: Taking the significant computational demand of TEMPO (lines 307-308) into account, I think the demonstrated performance improvement is marginal.**\n\n**A5:** The code is much faster now! Please refer to A7 to Reviewer 9zSq."}, {"review_id": "xUGHs9gAsu", "reviewer": "Reviewer_RGcg", "summary": "The authors present an elegant approach to combine meta-learning and world models. The authors cover the components of their Task-aware Environment Modeling Pipeline with Bi-level Optimization (TEMPO) approach. They then cover experiments across continuous and discrete domains, finding that TEMPO meets or exceeds the performance of DreamerV2. In addition, they present results of ablations to demonstrate the value of their specific approach. \n\nEdit: I appreciate the author's rebuttal and further discussion. I was already positive on the paper and remain so. I did not further increase my rating from an accept as there are issues that cannot be addressed by a revision.", "questions": "1. Why did the authors choose DreamerV2 as the one baseline?\n2. What do the authors think is happening with Hopper Stand?", "limitations": "No negative societal impacts likely, limitations acknowledged.", "rating": 7, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "This is a very strong paper. It is well-organized, well-written, and well-motivated. The authors identify a timely opportunity to combine meta-learning and world models to gain the values of both. It's perhaps a relatively simple idea, but the execution is at a very high level. The quantity and quality of experiments are also excellent as is the discussion putting those results into context. This is the type of paper I immediately want to share with my students, so it's just a shame I have to wait for it to be accepted!", "weaknesses": "The strengths aside, there are a few weaknesses that could be addressed to improve the paper. These keep me from advocating for a higher score than accept (for now). First, while DreamerV2 is a strong baseline, there have been other world model approaches released after it with improved performance. Notably, DreamerV3. The inclusion of these other baselines or at least discussion around these other potential baselines could have benefitted the paper and better justified the choice of DreamerV2. Secondly and relatedly, the related work (Section 4) currently reads more of a list of historical events without any context for how the authors view these works in relation to their own. This is particularly bad at the end, when approaches are just individually described, but no reason given as to why. Finally, there's some interesting results that would have benefitted from more discussion. Specifically, whatever was happening with Hopper Stand. In both Figure 2 and Figure 5 its a clear outlier. Some discussion as to what is happening here would improve the paper, even just in the appendix.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Dear Reviewer RGcg,\n\nWe sincerely thank you for giving us encouraging comments and recognizing the contribution of our work. Here are our answers to your questions:\n\n---\n**Q1: While DreamerV2 is a strong baseline, there have been other world model approaches released after it with improved performance. Notably, DreamerV3. The inclusion of these other baselines or at least discussion around these other potential baselines could have benefitted the paper and better justified the choice of DreamerV2. Why did the authors choose DreamerV2 as the one baseline?**\n\n**A1:** The Dreamer series is one of the most, if not the most, recognized model-based RL algorithms in recent years. Its biggest strength, as we see it, is its generality. It provides state-of-the-art performance in a variety of envs and tasks right out of the box, while being intuitive in design. Perhaps the only other method that can rival the fame and generality of Dreamer is MuZero, although, to our knowledge, no one has pitted them together in benchmarks like DMC and Atari under identical env settings. For TEMPO, we want to inherit the generality of Dreamer and further boost its performance by incorporating value awareness. Given that we were in a time rush for this work, we thought the easiest way to effectively support our method is to set Dreamer as the main baseline. However, we do believe that comparing TEMPO with other methods like MuZero will provide valuable insights into TEMPO, as well as Dreamer.\n\nAs for why build TEMPO on top of DreamerV2 rather than DreamerV3, the reason is quite practical. When we found out about DreamerV3, although it has a significant improvement over DreamerV2 in terms of generality, the official implementation is written in Google JAX, rather than PyTorch or TensorFlow, which we have experience with. By that time, we have already done part of our work based on DreamerV2. With limited time, we decided to stick with v2. Then again, for Dreamer, TEMPO is an additive bonus boost algorithm that optimizes the data usage, it is absolutely compatible with v3. We are excited to see TEMPO improve DreamerV3 in the future!\n\n---\n**Q2: The related work (Section 4) currently reads more of a list of historical events without any context for how the authors view these works in relation to their own. This is particularly bad at the end, when approaches are just individually described, but no reason given as to why.**\n\n**A2:** This is a great suggestion! We will try to polish the related work section and further clarify their relationship to our work in the revised version if we can manage to fit the page limit. For now, readers can still get a general and intuitive idea of the relationship between our work and others in the introduction section.\n\n---\n**Q3: There's some interesting results that would have benefitted from more discussion. Specifically, whatever was happening with Hopper Stand. Some discussion as to what is happening here would improve the paper, even just in the appendix. What do the authors think is happening with Hopper Stand?**\n\n**A3:** Again, we thank you for this great suggestion! We will add further discussion of our results to the appendix in the revised version.\n\nAs we pointed out in A6 to Reviewer 9zSq, we believe the distance between naive state representation from MLE and the task-aware state representation can vary on different tasks. When the distance is large, meaning that the features that can effectively lower reconstruction error don't coincide well with the task-specific features, TEMPO can have a significant advantage (e.g. Finger Spin). We think Hopper Stand is one of those tasks where such distance is small in the first place, and the RSSM itself can capture the task-specific features well enough, so TEMPO shows less advantage."}, {"review_id": "5GtUAjLxEf", "reviewer": "Reviewer_YURA", "summary": "The authors propose a novel task-aware model learning scheme for model-based RL, which combines the idea of value aware losses with bilevel optimization to obtain a stable training objective from SOTA pixel-based RL approaches such as Dreamer. The method uses the VAML loss as a surrogate objective for learning a weighing function which learns to provide weights to the samples of the ELBO training objective of Dreamer. The loss is therefore able to put more emphasis on samples in which the ELBO needs to be minimized more accurately to prevent large value function differences. The authors evaluate their method with empirical comparisons on some Atari games and DMC environments.", "questions": "For the naive weighing approach: it is unclear from the paper what exactly is gained from the learned weighting over the naive one, especially since the loss seems to be optimized by the naive weight for some mild assumptions. Therefore a qualitative experiment would be very helpful to properly compare: is the learned weight more smooth, are there generalization effects that it is able to take advantage of, etc.?\n\nThe authors should discuss the limitations of the proposed approach. Compared to other, more granular approaches such as Vagram (https://openreview.net/forum?id=4-D6CZkRXxI) and TIA (https://arxiv.org/abs/2106.15612), the per-sample weights would not be able to differentiate \"RL relevant\" dimensions from distracting ones, correct? This should be mentioned and briefly discussed as an important trade-off.", "limitations": "The authors describe the limitations shortly with a focus on computational efficiency. For the domain of the paper, this seems sufficient.", "rating": 7, "confidence": 5, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "The paper proposes a novel approach to task-aware model learning that deals well with previous established problems in the approach and extends well to the more complicated RSSM structures commonly used in pixel-based control.\n\nThe paper is overall clearly written without unnecessary notational baggage (see below for improvement suggestions).\n\nThe approach seems to offer some improvements in a selection of environments.\n\nThe proposed methodology could be extended more, especially with additional experiments clarifying how the learned network improves over the naive weighing strategy, which could make this an important step for future research into the area. This is especially promising since the paper elegantly sidesteps documented problems with the VAML approach.", "weaknesses": "The authors (correctly) note that VAML (as is) is inapplicable in a POMDP setting, but propose a method that relies on inferring value functions from the learned \"prior\" in the RSSM, which by the same logic, should not be doable. Some discussion of this would be good.\n\nSince the paper discusses task aware losses, comparisons to MuZero/EfficientZero (for atari) and TD-MPC (for DMC) would be helpful in understanding the impact of the improvement. This is a non-critical improvement, but it would help contextualizing the results.\n\nThe paper does not show whether the loss truly improves (in a VAML sense) over Dreamer. An additional Figure that shows whether the new loss truly results in a lower VAML error and whether this correlates with improved performance would be great.\n\n164: Notationally, the loss $\\mathcal{L}-\\text{upper}$ is independent of $\\alpha$, which is slightly strange. The equations could be rewritten for more clarity, i.e. by moving up equation (8). To make this easier for readers unfamiliar with bilevel optimization, it could also help to reduce the exposition on the RSSM, as most of the details aren't needed for understanding this work and can easily be gathered from the prior work by Hafner, and instead dedicate the space to explaining the bi-level optimization in a bit more detail.\n\n171: It is unclear to me why normalizing the value functions would be sufficient to prevent all weights from becoming small. Are weights also normalized during training? Otherwise wouldn't the algorithm still be incentivized to output 0 weights? Maybe I misunderstood the described problem (as the outer objective does not include the weights, why would it incentivize minimization?).\n\n196: sigmoid-ed -> a sigmoid function is applied.\n\nFigure 2: the chosen tasks are all from the \"easy\" and \"medium\" tasks (including those in the appendix) (classification by Yarats DrQ). Why was the suite sub-sampled in this way? Evaluations on hard tasks would be helpful, especially since most environments don't show marked difference between TEMPO and Dreamer. The same question of game selection is also applicable to the Atari benchmark.\n\nFigure 3: why are 6 games shown and only curves for 3?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Dear Reviewer YURA,\n\nWe sincerely appreciate your detailed feedback and constructive suggestions! Here is our response to your concerns:\n\n---\n**Q1: The authors (correctly) note that VAML (as is) is inapplicable in a POMDP setting, but propose a method that relies on inferring value functions from the learned \"prior\" in the RSSM, which by the same logic, should not be doable.**\n\n**A1:** The main reason for VAML not being straightforwardly applicable in general POMDP settings is that state space is not predefined. One could still use VAML by simply changing the states to observations, and learning the model and value function to work with the original observation space, but doing so is a waste of compact latent space built by RSSM.\n\nThe motivation for our modified loss is that, in any VAE-like model, the latent sampled from the prior distribution is considered a predictive latent as the name \"prior\" indicates, while the latent sampled from the posterior distribution is considered more a \"ground-truth\" latent as it is inferred from encoded observation. Since the latent space, which is a learned state space, is totally accessible, we propose to modify VAML by replacing the ground truth environmental states with posterior latents, and predicted states from the model with prior latents. The new loss is to evaluate RSSM's performance in the latent space it built, rather than predefined state space like VAML.\n\n---\n**Q2: Since the paper discusses task aware losses, comparisons to MuZero/EfficientZero (for atari) and TD-MPC (for DMC) would be helpful in understanding the impact of the improvement. This is a non-critical improvement, but it would help contextualizing the results.**\n\n**A2:** More comparisons can definitely be helpful! Please refer to A1 to Reviewer RGcg.\n\n---\n**Q3: The paper does not show whether the loss truly improves (in a VAML sense) over Dreamer. An additional Figure that shows whether the new loss truly results in a lower VAML error and whether this correlates with improved performance would be great.**\n\n**A3:** As we pointed out in A1, our new loss is essentially VAML loss in the latent space of RSSM. We added a new figure (a modified version of the current appendix E) comparing DreamerV2's training curve with our TEMPO, which can be found in Figure 2 of the rebuttal pdf.\n\n---\n**Q4: It is unclear to me why normalizing the value functions would be sufficient to prevent all weights from becoming small.**\n\n**A4:** This is a great question! First, a small mistake on our behalf. We weren't being totally informative by writing \"The meta weights are sigmoid-ed after the output layer\". A bias is also added after the sigmoid so that the weights center around 1, i.e. $\\text{weight} = \\sigma (\\text{output}) + 0.5$ so that the weights range from 0.5 to 1.5, or $\\text{weight} = 0.5 \\times \\sigma (\\text{output}) + 0.75$ so that the weights range from 0.75 to 1.25. We tried both in our initial experiments and didn't observe a noticeable difference, so we stuck with the latter in the hope of stable training. We will add this information in the revised version.\n\nBack to the question, without the normalization on the values, the task-aware loss actually rises as training continues, even though the meta weighter is trained to minimize it. This makes sense as the scale of values rises as the agent gets better at the task. However, we observe that most of the meta weights shrink to near the minimal value (i.e. 0.75 in our setting) after some training. We suspect this is because the weighter wants to stop the model from learning environment dynamics, so that the agent can't derive a good task policy from the model, therefore stopping the values from rising. Normalizing the values helps with this problem greatly as it allows the weighter to focus on minimizing the loss within the current value scale.\n\n---\n**Q5: How were the tasks chosen?**\n\n**A5:** Our concern when choosing tasks is mainly the computational cost. Since we want to run experiments on as many tasks as possible to demonstrate TEMPO's novelty, we couldn't afford to run 5e6 steps for DMC and 2e8 steps for Atari as the Dreamer papers did. We only had a 1e6-step budget for DMC and a 1e7-step budget for Atari given limited time and computation resources. Therefore, for DMC, we focus on choosing tasks from different environments on which Dreamer can reach convergence in under 1e6 steps, and the same thing for Atari with 1e7 steps. We are definitely excited to see how TEMPO does on harder tasks in the future.\n\n---\n**Q6: Why are 6 games shown and only curves for 3?**\n\n**A6:** The rest are in the appendix due to the space limit. Please refer to A4 to Reviewer kk1X.\n\n---\n**Q7: For the naive weighing approach: it is unclear from the paper what exactly is gained from the learned weighting over the naive one. A qualitative experiment would be very helpful to properly compare the two.**\n\n**A7:** This is a great question! However, so far, we too are not sure how are the meta weights outperforming the naive weights. We have included a table of a few sets of naive weights and meta weights in Table 1 of the rebuttal pdf. The Pearson correlations between them show that the meta weights are generated in a non-trivial way. We intended to study this question further in the future.\n\n---\n**Q8: Compared to other, more granular approaches such as Vagram and TIA, the per-sample weights would not be able to differentiate \"RL relevant\" dimensions from distracting ones, correct? This should be mentioned and briefly discussed as an important trade-off.**\n\n**A8:** Thank you for this great suggestion! VaGraM, as a variant of VAML, does offer the advantage of value-aware model learning in separate state dimensions. It shares the same problem in applicability with VAML. However, similar modifications as what we have done to VAML can be made to VaGraM. The resulting loss can also be used in TEMPO's upper level. We will add more discussion of this in the revised appendix."}]

Author rebuttal:
Dear Reviewers,

Thank you for your comments and suggestions. We hope our response is sufficient to address your questions.

Please kindly let us know if you have any additional concerns.

Paper 3253 Authors



Meta-review: The paper proposes a task/value-aware model-based RL algorithm based on a bi-level optimization. It uses a Dreamer-like architecture, which solves an MLE criteria for model learning, as the low-level optimizer. In order to make this value-aware, it suggests using a VAML-like objective to train a network that weighs each term of the low-level optimization problem.

The consensus among reviewers is that the paper should be accepted. After reading the paper, I also agree. This is a good addition to the task/value/decision-aware model learning literature.

Reviewers provided several comments in how this paper can be improved. Specifically, I'd like to emphasize the followings:

- The number of seeds is 3, which is small. This is especially important for some results where there is a large overlap between figures, for example, Figure 5 in the main body, and Figures 1 and 2 in the appendix. I ask the authors to increase the seed or perform statistical significant tests.

- The contribution of this work and difference with prior work should be made more clear. The related work section is too compressed. I suggest having an expanded section in an appendix that helps position the paper properly.

- It might be worth trying different networks with varying level of capacity. Value-awareness shows its impact especially when the model is not capable to capture all aspects of the model. Given that TEMPO is only using a weighted MLE to tune the model, studying the effect of model capacity can be illuminating. (This is optional.)

- In my own reading of the paper, I feel initially confused by the use of s^post and s^prior in Eq. 4. This is defined in the paragraph after Eq. 1, but as far as I can see, they have not been explicitly defined. Clarification around this and the intuitive difference between these two states, would be helpful.
Decision: Accept (poster)