SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Huining Yuan Hongkun Dou Xingyu Jiang Yue Deng*
School of Astronautics, Beihang University, Beijing, China
{hnyuan, douhk, jxy33zrhd, ydeng}@buaa.edu.cn
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.
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 value-focused 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 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 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 \in [1:T]$ , high-dimensional observations  $o_{1:T}$  (usually images in visual-based cases), continuous or discrete vector-valued actions  $a_{1:T}$ , and scalar rewards  $r_{1:T}$ . The observations and rewards are generated by the black-box environment  $o_t, r_t \sim p(o_t, r_t | o_{<t}, a_{<t})$ , with actions generated by the agent  $a_t \sim p(a_t | o_{\leq t}, a_{<t})$ . The goal of RL agents is to maximize the expected sum of rewards  $\mathrm{E}_p(\sum_{t=1}^T 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 VAE (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.
Figure 1: Main workflow of the proposed bi-level model learning paradigm TEMPO, corresponding to steps 3-8 in Algorithm 1. We start with (a) inferring the states and meta weights with model  $M_{\theta}$ , value function  $V$ , and meta weighter  $N_{\alpha}$ , then sequentially perform (b) the lower optimization on the model and (c) the upper optimization on the meta weighter.
Consider a trajectory  $\tau$  in the environment, which comprises a sequence of observations, actions, and rewards  $\tau = \{o_t, a_t, r_t\}_{t=1}^T$ . 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
Deterministic state module:  $h_t = f_\theta(h_{t-1}, s_{t-1}, a_{t-1})$
Stochastic state module:  $s_t \sim p_\theta(s_t|h_t)$
Observation & reward module:  $o_t, r_t \sim p_\theta(o_t, r_t | h_t, s_t)$
Representation module:  $s_t \sim q_\theta(s_t | h_t, o_t)$
The deterministic state module recurrently outputs  $h_t$ , whereas the stochastic state module, which corresponds to the prior network in a conditional VAE, 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_\theta$  with parameters  $\theta$ .
All components of the model are trained to maximize a 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 inequality, the variational bound is written as (see the Appendix for the full derivation)
$$
\begin{array}{l} \log p _ {\theta} \left(o _ {1: T}, r _ {1: T} \mid a _ {1: T}\right) \geq \mathcal {L} ^ {\mathrm {M L E}} (\tau ; \theta) = \sum_ {t = 1} ^ {T} \mathrm {E L B O} _ {t} \left(o _ {\leq t}, r _ {\leq t}, a _ {<   t}; \theta\right) \\ = \sum_ {t = 1} ^ {T} \left(\underbrace {\mathrm {E} _ {q _ {\theta} \left(s _ {t} \mid h _ {t} , o _ {t}\right)} \left[ p _ {\theta} \left(o _ {t} , r _ {t} \mid s _ {t} , h _ {t}\right) \right]} _ {\text {r e c o n s t r u c t i o n}} - \underbrace {\mathrm {E} _ {q _ {\theta} \left(s _ {t - 1} \mid h _ {t - 1} , o _ {t - 1}\right)} \left[ \mathrm {K L} \left[ q _ {\theta} \left(s _ {t} \mid h _ {t} , o _ {t}\right) \| p _ {\theta} \left(s _ {t} \mid h _ {t}\right) \right] \right]} _ {\text {r e g u l a r i z a t i o n}}\right) \tag {2} \\ \end{array}
$$
Here, we denote the objective function as  $\mathcal{L}^{\mathrm{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  $\epsilon$  from the ground truth  $s^{\prime}$ , the state value predicted by the critic will also have an error  $V(s^{\prime} + \epsilon) - V(s^{\prime})$ . 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  $\hat{p}$ , a distribution over the state-action space  $\mu$ , and a value function  $V$ , the VAML loss is written as
$$
\mathcal {L} ^ {\mathrm {V A M L}} (p, \hat {p}; \mu , V) = \int \mu (s, a) \left| \underbrace {\int p \left(s ^ {\prime} \mid s , a\right) V \left(s ^ {\prime}\right) \mathrm {d} s ^ {\prime}} _ {\text {e n v i r o n m e n t v a l u e e s t i m a t e}} - \underbrace {\int \hat {p} \left(s ^ {\prime} \mid s , a\right) V \left(s ^ {\prime}\right) \mathrm {d} s ^ {\prime}} _ {\text {m o d e l v a l u e e s t i m a t e}} \right| ^ {2} \mathrm {d} (s, a) \tag {3}
$$
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, 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/integrals is challenging in a realistic RL environment, where the state space is often vast and continuous. Additionally, VAML 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  $\tau$ , resulting in a parsimonious task-aware loss function, which we name V-VAML (Variational VAML)
$$
\mathcal {L} ^ {\mathrm {V - V A M L}} (\tau ; V, \theta) = \sum_ {t = 1} ^ {T} \left| V \left(s _ {t} ^ {\text {p o s t}}\right) - V \left(s _ {t} ^ {\text {p r i o r}}\right) \right| ^ {2} \tag {4}
$$
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  $\theta$  deterministically through reparameterization, i.e.  $s_t^{\mathrm{post}}(o_{\leq t},a_{< t};\theta)$  and  $s_t^{\mathrm{prior}}(o_{< t},a_{< t};\theta)$ . 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 a dataset. It is also evident that  $\mathcal{L}^{\mathrm{V - VAML}}$  has some similarities with the KL regularizer in  $\mathcal{L}^{\mathrm{MLE}}$ , where a large distance between the posterior and prior is undesirable in both objectives. The difference is that  $\mathcal{L}^{\mathrm{V - VAML}}$  evaluates the impact of such distance using the value disparity. One could try somehow replacing the KL regularizer with  $\mathcal{L}^{\mathrm{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
Algorithm 1: Task-aware Environment Modeling Pipeline with bi-level Optimization (TEMPO)
Input: Environmental trajectory  $\tau$ , world model  $M_{\theta}$ , meta weighter  $N_{\alpha}$ , value function update steps  $K$ , learning rate for model  $\eta$ , learning rate for meta weighter  $\lambda$   
1 while  $\theta$  or  $\alpha$  not converged do  
2 for update step  $k = 1 \ldots K$  do  
3 Infer states with world model  $h_{1:T}$ ,  $s_{1:T}^{post}$ ,  $s_{1:T}^{prior} = M_{\theta}(\tau)$ ;  
4 Compute values  $\{V(s_t^{post}), V(s_t^{prior})\}_{t=1}^T$ ;  
5Normalize values  $\{V(s_t^{post}), V(s_t^{prior})\}_{t=1}^T \leftarrow \text{Norm}(\{V(s_t^{post}), V(s_t^{prior})\}_{t=1}^T)$ ;  
6 Compute meta weights  $w_{1:T} = \{N_{\alpha}(h_t, s_t^{post}, s_t^{prior}, V(s_t^{post}), V(s_t^{prior}))\}_{t=1}^T$ ;  
7 Update world model  $\theta \gets \theta + \eta \cdot \frac{\partial \mathcal{L}^{\text{lower}}(\tau; w_{1:T}, \theta)}{\partial \theta}$ ;  
8 Update meta weighter  $\alpha \gets \alpha - \lambda \cdot \frac{\partial \mathcal{L}^{\text{upper}}(\tau; V, \theta)}{\partial \theta} \frac{\partial^2 \mathcal{L}^{\text{lower}}(\tau; w_{1:T}, \theta)}{\partial \theta \partial \alpha}$ ;  
9 end  
10 end
importance weight in a task-aware fashion. To achieve this, we introduce an additional meta-weight network  $N_{\alpha}$  (with parameters  $\alpha$ ) that outputs a meta weight  $w_{t}$  for each  $\mathrm{ELBO}_t(o_{\leq t}, r_{\leq t}, a_{<t}; \theta)$  in  $\mathcal{L}^{\mathrm{MLE}}$ , based on the current states and values (Figure 1(a))
$$
w _ {t} = N _ {\alpha} \left(h _ {t}, s _ {t} ^ {\text {p o s t}}, s _ {t} ^ {\text {p r i o r}}, V \left(s _ {t} ^ {\text {p o s t}}\right), V \left(s _ {t} ^ {\text {p r i o r}}\right)\right) \tag {5}
$$
With these meta weights, the model objective changes from  $\mathcal{L}^{\mathrm{MLE}}$  to a weighted sum of  $\mathrm{ELBO}_t\mathrm{s}$ . The meta weights as well as the meta weighter are trained to minimize our task-aware model loss  $\mathcal{L}^{\mathrm{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)
$$
\begin{array}{l} \min  _ {\alpha} \mathcal {L} ^ {\text {u p p e r}} (\tau ; V, \theta^ {*}) = \mathcal {L} ^ {\text {V - V A M L}} (\tau ; V, \theta^ {*}) = \sum_ {t = 1} ^ {T} \left| V \left(s _ {t} ^ {\text {p o s t}} \left(o _ {\leq t}, a _ {<   t}; \theta^ {*}\right)\right) - V \left(s _ {t} ^ {\text {p i r o r}} \left(o _ {<   t}, a _ {<   t}; \theta^ {*}\right)\right) \right| ^ {2} \\ \text {s . t .} \theta^ {*} = \underset {\theta} {\arg \max } \mathcal {L} ^ {\text {l o w e r}} (\tau ; w _ {1: T}, \theta) = \underset {\theta} {\arg \max } \sum_ {t = 1} ^ {T} w _ {t} \cdot \operatorname {E L B O} _ {t} \left(o _ {\leq t}, r _ {\leq t}, a _ {<   t}; \theta\right) \tag {6} \\ \end{array}
$$
Notice that  $\theta^{*}$  depend on meta weights  $w_{1:T}$  and parameter  $\alpha$  through the gradients of  $\mathcal{L}^{\mathrm{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  $\mathcal{L}^{\mathrm{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_t^{\mathrm{post}}), V(s_t^{\mathrm{prior}})\}_{t=1}^T$  before computing  $\mathcal{L}^{\mathrm{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  $\theta$  through a single gradient update (Figure 1(b))
$$
\theta^ {\prime} = \theta + \eta \cdot \nabla_ {\theta} \mathcal {L} ^ {\text {l o w e r}} (\tau ; w _ {1: T}, \theta) \tag {7}
$$
where  $\eta$  denotes the learning rate for the model. Then, we perform the upper-level optimization on the meta-weighter parameters  $\alpha$ , for which we can simply calculate  $\mathcal{L}^{\mathrm{upper}}$  using  $\theta'$  and take the gradient w.r.t  $\alpha$
Cartpole Swingup
Figure 2: Evaluation of TEMPO on continuous control tasks from DeepMind Control Suite. (a), a graphical demonstration of the 9 environments used in the evaluation, from left to right: Acrobot, Cartpole, Cheetah, Cup, Finger, Hopper, Pendulum, Quadruped, and Walker. (b), the evaluation curves of 3 seeds and their average significance, with lines showing the mean scores, and shaded areas showing the standard deviations.
$$
\begin{array}{l} \nabla_ {\alpha} \mathcal {L} ^ {\text {u p p e r}} (\tau ; V, \theta^ {\prime}) = \frac {\partial \mathcal {L} ^ {\text {u p p e r}} (\tau ; V , \theta^ {\prime})}{\partial \theta^ {\prime}} \frac {\partial \theta^ {\prime}}{\partial \alpha} \\ = \frac {\partial \mathcal {L} ^ {\text {u p p e r}} (\tau ; V , \theta^ {\prime})}{\partial \theta^ {\prime}} \frac {\partial^ {2} \mathcal {L} ^ {\text {l o w e r}} (\tau ; w _ {1 : T} , \theta)}{\partial \theta \partial \alpha} \tag {8} \\ \end{array}
$$
However, the problem with this direct approach is that we need to infer the states (i.e.  $h_{1:T}$ ,  $s_{1:T}^{\mathrm{post}}$ ,  $s_{1:T}^{\mathrm{prior}}$ ) twice, one time for each objective, for they come from different model parameters (i.e.  $\theta$  and  $\theta'$ ). Therefore, to cut down the computational cost, we propose to empirically approximate  $\theta'$  with  $\theta$  since they are only one update away. This results in an elegant update formula for  $\alpha$  (Figure 1(c))
$$
\begin{array}{l} \alpha^ {\prime} = \alpha - \lambda \cdot \nabla_ {\alpha} \mathcal {L} ^ {\text {u p p e r}} (\tau ; V, \theta^ {\prime}) \\ \approx \alpha - \lambda \cdot \frac {\partial \mathcal {L} ^ {\text {u p p e r}} (\tau ; V , \theta)}{\partial \theta} \frac {\partial^ {2} \mathcal {L} ^ {\text {l o w e r}} (\tau ; w _ {1 : T} , \theta)}{\partial \theta \partial \alpha} \tag {9} \\ \end{array}
$$
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.