SlowGuess/NeurIPS_2023_OCR huggingface dataset:
We implement our TEMPO framework on top of DreamerV2 using the official implementation of Hafner et al. (2020). Specifically, we build the meta weightter 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 \times \sigma(\text{output}) + 0.75$ . An Adam optimizer (Kingma and Ba, 2014) with a learning rate of  $1\mathrm{e} - 4$  is used for updating the meta weightter. 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.
Figure 3: Evaluation of 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 of 3 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 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.
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 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.
Figure 4: Training curve TEMPO on Cartpole Swingup from DeepMind Control Suite. The lines show the means and shaded areas show the standard deviations of 3 seeds.
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
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 \times 64 \times 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 \times 64 \times 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.
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, 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_t^{\mathrm{post}}$  and  $s_t^{\mathrm{prior}}$ , and state values  $V(s_t^{\mathrm{post}})$  and  $V(s_t^{\mathrm{prior}})$  (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.
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-VAE (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. VAML 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.
This work was supported in part by the National Natural Science Foundation of China under Grant 62325101 and Grant 62031001.
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. 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-Arnold, 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 Abdelmaleki, 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 Schon, 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 variances 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  $= \sigma (output) + 0.5$  and weight  $= 0.5 \times \sigma (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 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. 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 perform 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.
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}) = \prod_{t=1}^{T}q(s_t|h_t,o_t)$ , is written as
$$
\begin{array}{l} \log p \left(o _ {1: T}, r _ {1: T} \mid a _ {1: T}\right) \\ = \log \int \prod_ {t = 1} ^ {T} p \left(s _ {t} \mid h _ {t}\right) \cdot p \left(o _ {t}, r _ {t} \mid s _ {t}, h _ {t}\right) d s _ {1: T} \\ = \log \int \prod_ {t = 1} ^ {T} q \left(s _ {t} \mid h _ {t}, o _ {t}\right) \cdot \frac {p \left(s _ {t} \mid h _ {t}\right) p \left(o _ {t} , r _ {t} \mid s _ {t} , h _ {t}\right)}{q \left(s _ {t} \mid h _ {t} , o _ {t}\right)} d s _ {1: T} \\ = \log \mathrm {E} _ {q \left(s _ {1: T} \mid o _ {1: T}, a _ {1: T}\right)} \left[ \prod_ {t = 1} ^ {T} \frac {p \left(s _ {t} \mid h _ {t}\right) p \left(o _ {t} , r _ {t} \mid s _ {t} , h _ {t}\right)}{q \left(s _ {t} \mid h _ {t} , o _ {t}\right)} \right] \tag {10} \\ \geq \mathrm {E} _ {q (s _ {1: T} | o _ {1: T}, a _ {1: T})} \left[ \sum_ {t = 1} ^ {T} \log p \left(o _ {t}, r _ {t} | s _ {t}, h _ {t}\right) + \log p \left(s _ {t} | h _ {t}\right) - \log q \left(s _ {t} | h _ {t}, o _ {t}\right) \right] \\ = \sum_ {t = 1} ^ {T} \left(\mathrm {E} _ {q \left(s _ {t} \mid h _ {t}, o _ {t}\right)} \left[ p \left(o _ {t}, r _ {t} \mid s _ {t}, h _ {t}\right) \right] - \mathrm {E} _ {q \left(s _ {t - 1} \mid h _ {t - 1}, o _ {t - 1}\right)} \left[ \mathrm {K L} \left[ q \left(s _ {t} \mid h _ {t}, o _ {t}\right) \| p \left(s _ {t} \mid h _ {t}\right) \right] \right]\right) \\ \end{array}
$$
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 λ-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	1e-4
Meta weighter	Number of dense layers	5
Hidden dimension	400
Activation function	ELU
Normalization	Batch
Learning rate	1e-4
for DMC	Input dimension	200 + 2 × 32 + 1 × 2
for Atari	Input dimension	600 + 2 × 32 × 32 + 1 × 2
Common	Optimizer	Adam
Gradient clipping	100
Adam epsilon	1e-5
Weight decay	1e-6
Table 1: Main hyperparameters of TEMPO. DMC stands for DeepMind Control tasks. We use the default setting of Dreamerv2 in World model, Agent, and Common. All experiments were run on a single Nvidia RTX 3090 GPU with Python 3.7 and Tensorflow 2.6. Refer to our sample code for full implementation details.
D Full results on continuous control tasks
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.
E Full results on discrete control tasks
Figure 8: Evaluation of TEMPO on discrete control tasks from Atari video games. The lines show mean scores and the shaded areas show the standard deviation across 3 random seeds.
F Training curves
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 3 random seeds.
G Sample weights