Title: Task-aware world model learning with meta weighting via bi-level optimization
Venue: NeurIPS 2023
Accepted: yes
Keywords: Model-based reinforcement learning

world model

generative model

meta-learning

bi-level optimization
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

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.

1 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, or computationally expensive simulators (Moerland et al., 2023).

Model-based reinforcement learning (MBRL) aims to enhance the sample efficiency and generalization 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 to provide the agent with the ability to generate simulated experiences or predict the outcome of actions. In the behavior learning stage, the agent learns its policy by interacting with the model without having to take actions in the real environment. This paradigm has received a lot of attention, and significant progress has been made in both model learning and behavior learning. (Sutton, 1991; Ha and Schmidhuber, 2018; Janner et al., 2019; Kaiser et al., 2019; Hafner et al., 2019a; Schrittwieser et al., 2020).

One of the most common approaches to building a model is to learn a deep generative model through maximum likelihood estimation (MLE) on environmental trajectories (Buesing et al., 2018; Hafner et 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-art model-free RL agents for task behavior (Ha and Schmidhuber, 2018; Kaiser et al., 2019). However, MLE reconstructs all information from

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 is 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 to 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 et 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-tailored planning algorithms for behavior learning (Schrittwieser 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 valuefocused learning can be useful for learning an effective policy.

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 a proposed task-aware model loss. Under this hierarchical framework, the meta weighter in the upper level of optimization learns to prioritize those samples with a 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 a 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 DreamerV2 (Hafner et al., 2020), a powerful MBRLalgorithm 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 RL 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-the-art 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 demonstrate the advantage of our proposed meta-weighting mechanism.

2 Task-aware 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 t ∈ [1 : T ] , highdimensional observations o 1: T (usually images in visual-based cases), continuous or discrete vectorvalued actions a 1: T , and scalar rewards r 1: T . The observations and rewards are generated by the blackbox environment o t , r t ∼ p ( o t , r t | o <t , a <t ) , with actions generated by the agent a t ∼ p ( a t | o ≤ t , a <t ) . The goal of RL agents is to maximize the expected sum of rewards E p ( ∑ T t =1 r t ) .

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 a MLE-based world model. Finally, we introduce our main contribution, a bi-level model learning framework where a meta-weighting 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 a probabilistic graphical model similar to a sequential V AE (Kingma and Welling, 2013; Sohn et al., 2015). It 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

M

States

V

Values

N

Meta

θ

weights

Lower optimization on world model: step 7

Netwoks&parameters

(b)

0Clower /00

θ

θ'

Inputs&outputs

Workflowofinference

Upper optimization on meta weighter: step 8

Workflow of

loweroptimization

(c)

Cupper /00

×

² Clower /00 0α

α

α

Workflow of

upperoptimization

Consider a trajectory τ in the environment, which comprises a sequence of observations, actions, and rewards τ = { o t , a t , r t } T t =1 . The RSSM introduces deterministic states h 1: T and stochastic states s 1: T as 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

The deterministic state module recurrently outputs h t , whereas the stochastic state module, which corresponds to the prior network in a conditional V AE, predicts the prior s t at each time step. These two state modules work together to perform the state transitions, while the observation & reward module acts as a decoder to reconstruct o t and r t from h t and s t . 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 s t . For clarity, we denote the entire model together as M θ with parameters θ .

All components of the model are trained to maximize a variational bound on the trajectory loglikelihood (also known as the Evidence Lower Bound, ELBO) using gradient-based methods (Kingma and Welling, 2013). With Jensen's inequality, the variational bound is written as (see the Appendix for the full derivation)

Here, we denote the objective function as L MLE . For each time step, the variational bound consists of two terms, i.e. the reconstruction accuracy of o t and r t , and the KL-divergence between the variational posterior and the predictive prior as a 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 to learn task behavior from simulated experiences in the compact latent state space.

2.2 The task-aware model loss

The RSSM is a 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 a given state-action pair ( s, a ) and a typical actor-critic RL agent. If the consequent state predicted by the model is off by ϵ from the ground truth s ′ , the state value predicted by the critic will also have an error V ( s ′ + ϵ ) -V ( s ′ ) . 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 a specific task, it is intuitive to evaluate a 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 p and its model approximation ˆ p , a distribution over the state-action space µ , and a value function V , the VAML loss is written as

The primary concept behind VAML is to penalize a 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, V AML 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/integrals is challenging in a realistic RL environment, where the state space is often vast and continuous. Additionally, V AML compels the model to conserve only the value-relevant characteristics, leading to the rejection of a considerable amount of semantic information contained in the environmental observations.

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 a parsimonious task-aware loss function, which we name V-VAML (Variational VAML)

Here, we introduce superscripts to differ posterior and prior states. Notice the states are outputs of the model, and depend on the model's parameters θ deterministically through reparameterization, i.e. s post t ( o ≤ t , a <t ; θ ) and s prior t ( o <t , a <t ; θ ) . We leave out the inputs for simplicity.

The intuition of V-V AML loss is quite similar to the original V AML loss, yet far easier to implement given a dataset. It is also evident that L V-VAML has some similarities with the KL regularizer in L MLE , where a large distance between the posterior and prior is undesirable in both objectives. The difference is that L V-VAML evaluates the impact of such distance using the value disparity. One could try somehow replacing the KL regularizer with L V-VAML 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 et al., 2018; Shu et al., 2019; Jiang et al., 2022), we propose to assign each training sample with different

importance weight in a task-aware fashion. To achieve this, we introduce an additional meta weighter network N α (with parameters α ) that outputs a meta weight w t for each ELBO t ( o ≤ t , r ≤ t , a <t ; θ ) in L MLE , based on the current states and values (Figure 1(a))

With these meta weights, the model objective changes from L MLE to a weighted sum of ELBO t s. The meta weights as well as the meta weighter are trained to minimize our task-aware model loss L V-VAML . This adds another level of optimization on top of the default MLE model learning, forming a bi-level framework with two hierarchical optimizations of orthogonal objectives, which we name Task-aware Environment Modeling Pipeline with bi-level Optimization (TEMPO)

Notice that θ ∗ depend on meta weights w 1: T and parameter α through the gradients of L lower . 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 to minimize L upper 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 a normalization with moving mean and variance on the state values { V ( s post t ) , V ( s prior t ) } T t =1 before computing L upper .

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 a single gradient update (Figure 1(b))

where η denotes the learning rate for the model. Then, we perform the upper-level optimization on the meta weighter parameters α , for which we can simply calculate L upper using θ ′ and take the gradient w.r.t α

(a)

t=3.322 p=0.080

t=15.985p=0.048

t=0.787 p=0.305

1000

1000

1000

Evaluation Return

800

800

800

600

600

600

(b)

400

400

400

200

200

200

0

0

0

0

2

4

6

8

10

0

2

4

6

8

10

0

2

4

6

8

10

EnvironmentalSteps

1e5

TEMPO

DreamerV2 －—D4PG(1e8）——A3C(1e8)

However, the problem with this direct approach is that we need to infer the states (i.e. h 1: T , s post 1: T , s prior 1: T ) twice, one time for each objective, for they come from different model parameters (i.e. θ and θ ′ ). Therefore, to cut down the computational cost, we propose to empirically approximate θ ′ with θ , since they are only one update away. This results in an elegant update formula for α (Figure 1(c))

It 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 to infer the states only once during an epoch of bi-level optimization. See Algorithm 1 for a pseudocode of our bi-level model-learning paradigm and Figure 1 for a graphical demonstration.

3 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 a 5-layer dense network (MLP) with concatenated states as input (see Equation 5) and scalar meta weight as output. All hidden dimensions of the network are set to 400. Batch normalization (Ioffe and Szegedy, 2015) and ELU (Clevert et al., 2015) activation are performed after each hidden layer. A 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 × σ ( output ) + 0 . 75 . An Adam optimizer (Kingma and Ba, 2014) with a 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.

10

10000

000000

(a)

Bank Heist

Kangaroo

Hero

t=3.291p=0.125

t = 1.598p = 0.146

Evaluation Return

1200

8000

14000

6000

800

10000

(b)

4000

6000

400

2000

2000

0

0

2

4

6

8

10

0

2

4

6

8

10

0

2

4

6

8

EnvironmentalSteps

1e6

TEMPO

-DreamerV2 ——IQN (1e7) ——Rainbow (1e7)

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 a number of state-of-the-art model-free RL agents including D4PG (Barth-Maron et al., 2018), A3C (Mnih et al., 2016), IQN (Dabney et al., 2018), and Rainbow (Hessel et al., 2018). For the model-free agents, we use the results reported by Tassa et al. (2018) and Castro et al. (2018). For DreamerV2, we use the default hyperparameters provided by Hafner et 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.

Wetuned 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 reach early convergence on a meaningless state representation. Empirically, TEMPO trains stably, and both objectives reach convergence after sufficient training (Figure 4). Refer to the Appendix for further detailed settings.

Following Henderson et al. (2018), we perform significance testing to examine TEMPO's superiority. In particular, we perform the one-tailed Welch's t -test (Welch, 1947) on TEMPO's and DreamerV2's results for every evaluation during the training process, and report the average t and p over all evaluations.

Furthermore, to verify the effectiveness and novelty of the proposed meta-weighting mechanism, we compare TEMPO's meta-weighting mechanism with naively weighting the training samples using the task-aware model losses. We also perform ablation studies on different inputs of the meta weighter network and different hyperparameters.

Upper

Lower

Obj.

Obj.

1e4

0.8

-1.08

0.4

-1.10

0.0

-1.12

0

1

2

3

4

5

Environmental Steps

1e5

Upperobjective(task aware)

Lower objective(MLE)

1000

1000

Evaluation Return

800

800

600

600

(a)

400

400

200

200

TEMPO

-Naive Weighting

0

0

0

2

4

6

8

10

0

2

4

6

8

10

?

Environmental Steps

1e5

1000

1000

Evaluation Return

800

800

600

600

(b)

400

400

-D+S+V(TEMPO)

-D+S

200

200

S

0

0

0

2

4

6

8

10

0

2

4

6

8

10

→

1e5

Environmental Steps

3.2 Continuous control

A subset of the DeepMind Control Suite with 9 tasks is used to evaluate TEMPO's performance in continuous control situations. The tasks are each from a different environment, as illustrated in Figure 2(a). Environmental observations are RGB images of shape 64 × 64 × 3 ; actions range from 1 to 12 dimensions; each episode starts with a randomized initial state and lasts for 1000 steps. We follow the protocol described by Hafner et al. (2020) and set the stochastic latent variable (i.e. s t ) of DreamerV2 and TEMPO to be a 32-dim continuous vector following Gaussian distribution.

We evaluate the TEMPO agent's environment return every 1e4 environmental steps, and compare the results with those of DreamerV2 and the final performance of D4PG and A3C after 1e8 steps. The results of 3 tasks are illustrated in Figure 2(b), where TEMPO achieved state-of-the-art 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 9 tasks.

3.3 Discrete control

We then evaluate TEMPO's performance on discrete control tasks with 6 Atari video games, as illustrated in Figure 3(a). The actions range from 3 to 18 dimensions. We render the environmental observations as gray-scale images of shape 64 × 64 × 1 , and set the stochastic latent variable (i.e. s t ) 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 1e7 environmental steps. We evaluate the TEMPO agent's environment return every 1e5 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.

F0001

800

Evaluation Return

600

400

200

0

0

2

4

6

8

10

Environmental Steps

1e5

DreamerV2

-TEMPO (hidden_dim = 200)

-TEMPO (activation =ReLU)

-TEMPO

-TEMPO (layers = 3)

-TEMPO (norm = LayerNorm)

3.4 Comparison with naive weighting strategy

In TEMPO, the meta weighter, a deep neural network, generates the meta weights based on state information. Given that the meta weighter is trained to minimize our task-aware model loss, one may naturally wonder: is it viable to naively use the losses themselves as 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 a training trajectory using a moving mean and variance. Subsequently, we add 1 to these normalized losses to 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 to 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 t , stochastic states s post t and s prior t , and state values V ( s post t ) and V ( s prior t ) (see Equation 5). To justify our design, we conduct an ablation study on the inputs to assess the impact of these inputs on the agent's performance. We sequentially remove the state values and deterministic state from the input, and compare the outcomes with our 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 hyperparameters

To examine the robustness of TEMPO's meta-weighting mechanism, we perform simple ablation studies on a 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 the face of reduced hidden dimensions and altered activation functions. While modification to the 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 to 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 et 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 a Dyna fashion (Sutton, 1991). For example, SimPLe (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 to 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 characterize 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 et al., 2019b,a, 2020; Ozair et al., 2021). Based on the RSSM, Dreamer (Hafner et al., 2019a) and DreamerV2 (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-V AE (Van Den Oord et al., 2017), which is then combined with Monte Carlo tree search (MCTS (Coulom, 2007)) for task behavior.

Observation-space models and state-space models can both be categorized as 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-end value prediction (Tamar et al., 2016; Silver et al., 2017; Farquhar et 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 a 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-based world models (lower level) and task-aware model loss (upper level). On one hand, one can find a suitable task-aware model loss to evaluate a MLE-based model (e.g. V AML for observation-space models and our V-VAML for state-space models), and build a bi-level optimization to train the model in a 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-VAML.

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 a novel task-aware loss function under a bi-level optimization. TEMPO achieves state-of-the-art performance on a variety of continuous and discrete control tasks, while demonstrating better asymptotic performance, training stability, and convergence speed. Ablation studies are performed to justify our 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, we observed that TEMPO trains around 40% slower and requires 80% more RAM than DreamerV2. Nonetheless, TEMPO contributes significant insights and opens up an exciting new avenue for environment modeling. We aspire that TEMPO will serve as a catalyst for novel ideas and approaches for model-based RL in the future. A sample code of TEMPO is available at https://github.com/deng-ai-lab/TEMPO .

References

Ioannis Antonoglou, Julian Schrittwieser, Sherjil Ozair, Thomas K Hubert, and David Silver. Planning in stochastic environments with a learned model. In International Conference on 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 arXiv:1804.08617 , 2018.

G. Bellemare, Y. Naddaf, J. Veness, and M. Bowling. The arcade learning environment: An 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 arXiv:1802.03006 , 2018.

Pablo Samuel Castro, Subhodeep Moitra, Carles Gelada, Saurabh Kumar, and Marc G Bellemare. Dopamine: A research framework for deep reinforcement learning. arXiv preprint arXiv:1812.06110 , 2018.

Silvia Chiappa, Sébastien Racaniere, Daan Wierstra, and Shakir Mohamed. Recurrent environment simulators. arXiv preprint arXiv:1704.02254 , 2017.

Djork-Arné 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-massoud 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: Differentiable tree-structured models for deep reinforcement learning. arXiv preprint arXiv:1710.11417 , 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 2170-2179. 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: 5541-5552, 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, Ian 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 arXiv:2010.02193 , 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-ECCV 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 arXiv:1903.00374 , 2019.

Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.6980 , 2014.

Diederik P Kingma and Max Welling. Auto-encoding variational bayes. arXiv preprint arXiv:1312.6114 , 2013.

Nathan Lambert, Brandon Amos, Omry Yadan, and Roberto Calandra. Objective mismatch in model-based reinforcement learning. arXiv preprint arXiv:2002.04523 , 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 Broekens, Aske Plaat, Catholijn M Jonker, et al. Model-based reinforcement 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.02999 , 2018.

Evgenii Nikishin, Romina Abachi, Rishabh Agarwal, and Pierre-Luc Bacon. Control-oriented modelbased 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 in neural information processing systems , 28, 2015.

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-609, 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. Metaweight-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 DulacArnold, David Reichert, Neil Rabinowitz, Andre Barreto, et al. The predictron: End-to-end learning and planning. In International Conference on Machine 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 S 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, et 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 arXiv:2204.01464 , 2022.

Niklas Wahlström, Thomas B Schön, and Marc Peter Deisenroth. From pixels to torques: Policy learning with deep dynamical models. arXiv preprint arXiv:1502.02251 , 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.

Bernard L Welch. The generalization of 'student's'problem when several different population varlances are 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 arXiv:2006.10742 , 2020.

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 to 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, in the case of DMC, we focus on choosing tasks from different environments on which Dreamer can roughly reach convergence in under 1e6 steps. Similarly, for Atari, we focus on tasks that DreamerV2 could reach convergence within our 1e7 budget. We are certainly enthusiastic about exploring how TEMPO performs on 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 a similar MLP (Multi-Layer Perceptron) structure. Given that these modules can handle state representation well, it was a logical choice to build a 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 = σ ( output ) + 0 . 5 and weight = 0 . 5 × σ ( 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 we built TEMPO was to design a 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 TEMPOon 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 superior architecture and hyperparameters exist for TEMPO, and we remain committed to further exploration and refinement.

A.4 Performance analysis

Based on our findings, TEMPO demonstrates a 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 coincide well with the task-specific features, 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.

We also observe that TEMPO shows larger variances in some tasks than others. We suspect the reason is 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 multitasking 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.

B Derivations

The variational bound, i.e. Evidence Lower Bound (ELBO), for the Recurrent State-space Model (RSSM) with a variational posterior q ( s 1: T | o 1: T , a 1: T ) = ∏ T t =1 q ( s t | h t , o t ) , is written as

C Hyperparameters

D Full results on continuous control tasks

t= 0.940 p=0.371

t=3.322 p=0.080

t=0.567 p=0.371

009

1000

1000

500

800

800

400

600

600

300

400

400

200

100

200

200

0

0

0

0

2

4

6

8

10

0

2

4

6

8

10

0

2

4

6

8

10

Cup Catch

t=0.878 p=0.410

t=15.985p=0.048

t=0.787p=0.305

1000

1000

1000

Evaluation Return

800

800

800

600

600

600

400

400

400

200

200

200

0

0

0

0

1

2

3

4

5

0

2

4

6

8

10

0

2

4

6

8

10

Pendulum Swingup

WalkerWalk

t=1.004 p=0.425

t=0.519 p=0.381

t=0.661 p=0.327

1000

1000

1000

800

800

800

600

009

600

400

400

400

200

200

200

0

0

10

0

2

4

6

8

10

0

4

8

16

20

0

2

4

6

8

10

→

EnvironmentalSteps

1e5

TEMPO

(

E Full results on discrete control tasks

Bank Heist

Freeway

t=3.291 p=0.125

t= 0.666 p=0.368

t= 1.043 p=0.443

1200

1e4

12

30

800

20

8

400

10

Evaluation Return

0

0

0

2

4

6

8

10

0

2

4

6

8

10

0

1

2

3

4

Hero

Kangaroo

Pong

le3

t=1.598p =0.146

t= 1.046 p=0.258

t= 0.173 p=0.454

8000

14

20

10

6000

10

4000

0

6

-10

2000

-20

0

2

4

6

8

10

0

2

4

6

8

10

0

2

4

6

8

10

EnvironmentalSteps

1e6

TEMPO

F Training curves

Upper objective (task aware)

Lower objective (MLE)

1.0

-10800

0.8

-10900

0.6

-11000

0.4

-11100

Swingup

0.2

-11200

0.0

-11300

0

2

4

6

8

10

0

2

4

6

8

10

(a)

Cartpole S

Value

Unweighted ELBO

-11294

150

-11298

100

-11302

50

0

-11306

0

2

4

6

8

10

0

2

4

6

8

10

→

Environmental Steps

1e5

Upper objective (task aware)

Lower objective (MLE)

-3600

0.8

-3700

0.6

-3800

0.4

0.2

-3900

Bank Heist

0.0

-4000

0

2

4

6

8

10

0

2

4

6

8

10

(b)

Value

Unweighted ELBO

40

-3800

30

-3850

20

-3900

10

-3950

0

-4000

0

2

4

6

8

10

0

2

4

6

8

10

→

Environmental Steps

1e6

TEMPO

-DreamerV2

G Sample weights

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\nThe 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.\n\nWEAKNESSES:\nThe 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?\n\nQUESTIONS:\nFor 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.\n\nLIMITATIONS:\nThe authors describe the limitations shortly with a focus on computational efficiency. For the domain of the paper, this seems sufficient.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 5, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-04"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\nThis 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!\n\nWEAKNESSES:\nThe 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.\n\nQUESTIONS:\n1. Why did the authors choose DreamerV2 as the one baseline?\n2. What do the authors think is happening with Hopper Stand?\n\nLIMITATIONS:\nNo negative societal impacts likely, limitations acknowledged.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-07"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n- 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.\n\nWEAKNESSES:\nOne 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.\n\nQUESTIONS:\nN/A\n\nLIMITATIONS:\nN/A", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-07-09"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\n1. 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.\n\nWEAKNESSES:\n1. 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.\n\nQUESTIONS:\n1. 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.\n\nLIMITATIONS:\nThe 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.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-10"}