{'fraction_non_alphanumeric': 0.04418184223290268, 'fraction_numerical': 0.009700604083072446, 'mean_word_length': 4.88785046728972, 'pattern_counts': {'":': 0, '<': 6, '<?xml version=': 0, '>': 0, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
Dreaming: Model-based Reinforcement Learning by Latent Imagination without Reconstruction

In the present paper, we propose a decoder-free extension of Dreamer, a leading model-based reinforcement learning (MBRL) method from pixels. Dreamer is a sample- and cost-efficient solution to robot learning, as it is used to train latent state-space models based on a variational autoencoder and to conduct policy optimization by latent trajectory imagination. However, this autoencoding based approach often causes object vanishing, in which the autoencoder fails to perceives key objects for solving control tasks, and thus significantly limiting Dreamer's potential. This work aims to relieve this Dreamer's bottleneck and enhance its performance by means of removing the decoder. For this purpose, we firstly derive a likelihood-free and InfoMax objective of contrastive learning from the evidence lower bound of Dreamer. Secondly, we incorporate two components, (i) independent linear dynamics and (ii) the random crop data augmentation, to the learning scheme so as to improve the training performance. In comparison to Dreamer and other recent model-free reinforcement learning methods, our newly devised Dreamer with InfoMax and without generative decoder (Dreaming) achieves the best scores on 5 difficult simulated robotics tasks, in which Dreamer suffers from object vanishing.


Introduction
In the present paper, we focus on model-based reinforcement learning (MBRL) from pixels without complex reconstruction. MBRL is a promising technique to build controllers in a sample efficient manner, which trains forward dynamics models to predict future states and rewards for the purpose of planning and/or policy optimization. The recent study of MBRL in fully-observable environments [1,2,3,4] have achieved both sample efficiency and competitive performance with the state-of-the-art model-free reinforcement learning (MFRL) methods like soft-actor-critic (SAC) [5]. Although real-robot learning has been achieved with fully-observable MBRL [6,7,8,9,10,11], there has been an increasing demand for robot learning in partially observable environments in which only incomplete information (especially vision) is available. MBRL from pixels can be realized by introducing deep generative models based on autoencoding variational Bayes [12].
Object vanishing is a critical problem of the autoencoding based MBRL from pixels. Previous studies in this field [13,14,15,16,17,18,19] train autoencoding models along with latent dynamics models that generate imagined trajectories that are used for planning or policy optimization. However, the autoencoder often fails to perceive small objects in pixel space. The top part of Fig. 1 exemplifies this kind of failures where small (or thin) and important objects are not reconstructed in their correct positions. This shows the failure to successfully embed their information into the latent space, which significantly limits the training performance. This problem is a result of a loglikelihood objective (reconstruction loss) defined in the pixel space. Since the reconstruction errors of small objects in the pixel space are insignificant compared to the errors of other objects and uninformative textures that occupy most parts of the image region, it is hard to train the encoder to perceive the small objects from the weak error signals. Also, we have to train the decoder which requires a high model capacity with massive parameters from the convolutional neural networks (CNN), although the trained models are not exploited both in the planning and policy optimization. Dreamer Dreaming Figure 1: The overview of the motivation and concept of this work. (Top) The concept of Dreamer's autoencoding-based representation learning [16], which often causes object vanishing as shown with the dashed-line ovals. The left three tasks are from DeepMind Control Suite [20], and the remaining two tasks are our original tasks which are assumed to represent industrial applications. (Bottom) The concept of Dreaming's representation learning, which trains a discriminator instead of the decoder to embed different samples to be spaced apart from each other. The learning scheme is characterized with the two key components; (i) Linear dynamics successfully constrains temporally consecutive samples are not distributed too further away. (ii) Image augmentation encourages only key features for control to be embedded into the latent space.
To avoid this complex reconstruction, some previous MBRL studies have proposed decoder-free representation learning [16,21] based on contrastive learning [22,23], which trains a discriminator instead of the decoder. The discriminator is trained by categorical cross-entropy optimization, which encourages latent embeddings to be sufficiently distinguishable among different embeddings. Nevertheless, to the best of our knowledge, no MBRL methods have achieved the state-of-the-art results on the difficult benchmark tasks of DeepMind Control Suite (DMC) [20] without reconstruction.
Motivated by these observations, this paper aims to achieve the state-of-the-art results with MBRL from pixels without reconstruction. This paper mainly focuses on the latest autoencoding-based MBRL method Dreamer, considering the success of a variety of control tasks (i.e.,DMC and Atari Games [24]), and tries to extend this method to be a decoder-free fashion. We adopt Dreamer's policy optimization without any form of modification. We call our extended Dreamer as Dreamer with InfoMax and without generative decoder (Dreaming). The concept of this proposed method is illustrated in the bottom part of Fig. 1. Our primary contributions are summarized as follows.
• We derive a likelihood-free (decoder-free) and InfoMax objective for contrastive learning by reformulating the variational evidence lower bound of the graphical model of the partially observable Markov decision process.
• We show that two key components, (i) an independent and linear forward dynamics, which is only utilized for contrastive learning, and (ii) appropriate data augmentation (i.e., random crop), are indispensable to achieve the state-of-the-art results.
In comparison to Dreamer and the recent cutting edge MFRL methods, Dreaming can achieve the state-of-the-art results on difficult simulated robotics tasks exhibited in Fig. 1 in which Dreamer suffers from object vanishing. The remainder of this paper is organized as follows. In Sec. 2, key differences from related work are discussed. In Sec. 3, we provide a brief review of Dreamer and contrastive learning. In Sec. 4, we first describe the proposed contrastive learning scheme in detail, and then introduce Dreaming. In Sec. 5, the effectiveness of Dreaming is demonstrated through simulated evaluations. Finally, Sec. 6 concludes this paper.

Related Work
Some of the most related work are contrastive predictive coding (CPC) [22] and contrastive forward model (CFM) [21]. Our work is highly inspired by CPC, and our contrastive learning scheme has similar components with CPC; e.g., a recurrent neural network and a bilinear similarity function. However, CPC has no action-conditioned dynamics models. Since CPC alone cannot generate imagined trajectories from arbitrary actions, CPC is only used as an auxiliary objective of MFRL. CFM heuristically introduces a similar decoder-free objective function like ours. A primary difference between ours and CFM is that CFM exploits a shared and non-linear forward model for both contrastive and behavior learning. As we discuss and demonstrate later in the following sections, this significantly limits the training performance. In addition, CPC and CFM do not introduce data augmentation. Furthermore, the relation between the evidence lower bound of time-series variational inference is not discussed in these literature.
MFRL methods with representation learning: A state-of-the-art MFRL method, contrastive unsupervised representation for reinforcement learning (CURL) [25], also makes use of contrastive learning with the random crop data augmentation. Deep bisimulation for control (DBC) [26] and discriminative particle filter reinforcement learning (DPFRL) [27] are other types of cutting edge MFRL methods, which utilize different concepts of representation learning without reconstruction.
MFRL methods without representation learning: Recently proposed MFRL methods, which include reinforcement learning with augmented data (RAD) [28], data-regularized Q (DrQ) [29], and simple unified framework for reinforcement learning using ensembles (SUNRISE) [30], have achieved state-of-the-art result without representation learning. All these work employ the random crop data augmentation as an important component of their method. Let us begin by considering the graphical model illustrated in Fig. 2, whose joint distribution is defined as follows:

Autoencoding Variational Bayes for Time-series
where z, x, and a denote latent state, observation, and action, respectively. As in the case of well-known variational autoencoders (VAEs) [12], generative models p(z t+1 |z t , a t ), p(x t |z t ) and inference model q(z t |x ≤t , a <t ) can be trained by maximizing the evidence lower bound [15]: If the models are defined to be differentiable and trainable, this objective can be maximized by the stochastic gradient ascent via backpropagation.
Multi-step variational inference is proposed in [15] to improve long-term predictions. This inference, named latent overshooting, involves the multi-step objective J KL k defined as: where p(z t |z t−k , a <t ) : For the purpose of planning or policy optimization, not only for the dynamics model p(z t+1 |z t , a t ) but also the reward function p(r t |z t ) is also required. To do this, we can simply regard the rewards as observations and learn the reward function p(r t |z t ) along with the decoder p(x t |z t ). For readability, we omit the specifications of the reward function p(r t |z t ) in the following discussion. Although we remove the decoder later, the reward function and its likelihood objective are kept untouched.

Recurrent State Space Model and Dreamer
The recurrent state space model (RSSM) is a latent dynamics model equipped with an expressive recurrent neural network, realizing accurate long-term prediction. RSSM is used as an essential component of various MBRL methods from pixels [15,16,17,19,31] including Dreamer. RSSM assumes the latent z t comprises z t = (s t , h t ) where s t , h t are the probabilistic and deterministic variables, respectively. The generative and inference models of RSSM are defined as: Generative models : Inference model : where deterministic h t is considered to be the hidden state of the gated recurrent unit (GRU) f GRU (·) [32] so that historical information can be embedded into h t .
Dreamer [16] makes use of RSSM as a differentiable dynamics and efficiently learns the behaviors via backpropagation of Bellman errors estimated from imagined trajectories. Dreamer's training procedure is simply summarized as follows: (1) Train RSSM with a given dataset by optimizing Eq. (2).
(2) Train a policy from the latent imaginations.
(3) Execute the trained policy in a real environment and augment the dataset with the observed results. The above steps are iteratively executed until the policy performs as expected.

Contrastive Learning of RSSM
The original Dreamer paper [16] also introduced a likelihood-free objective by reformulating J likelihood of Eq. (2). By adding a constant log p(x t ) and applying Bayes' theorem, we get a decoder-free objective: where D batch denotes the mini-batch and the lower bound in the second line was from the Info-NCE (noise-contrastive estimator) mini-batch bound [33]. Let B be the batch size of D batch , J NCE is considered as a B-class categorical cross entropy objective to discriminate the positive pair (z t , x t ) among the other negative pairs (z t , x ( = x t )). In this interpretation, p(z t |x t ) can be considered as a discriminator to discern the positive pairs. Representation learning with this type of objective is known as contrastive learning [22,23] that encourages the embeddings to be sufficiently seperated from each other in the latent space. However, the experiment in [16] has demonstrated that this objective significantly degrades the performance compared to the original objective J likelihood .

Deriving Another Contrastive Objective
We propose to further reformulate J NCE of Eq. (5) by introducing a multi-step prediction model: The accent ofp implies that an independent dynamics model from p(z t |z t−1 , a t−1 ) in Eq. (2) can be employed here. By multiplying a constant E q(z t−k |·) [p(z t |z t−k , a <t )/p(z t |z t−k , a <t )] = 1, we obtain an importance sampling form of J NCE as: For computational simplicity, we approximate the likelihood ratio q(z t |·)/p(z t |·) as a constant and assume that the summation of J NCE across batch and time dimension is approximated as: batch,time We further import the concept of overshooting and optimize J NCE k along with J KL k on multi-step prediction of varying ks. Finally, the objective we use to train RSSM is: Note that J NCE k and J KL k have different dynamics model (i.e.,p(z t |·) and p(z t |·), respectively).

Relation among the Objectives
As shown in Appx. A, J NCE is a lower bound of the mutual information I(x t ; z t ), while J NCE k is a bound of I(x t ; z t−k ). Since the latent state sequence is Markovian, we have the data processing inequality as I(x t ; z t ) ≥ I(x t ; z t−k ). In other words, J NCE and approximately derived J NCE k share the same InfoMax upper bound metrics. An intuitive motivation to introduce J NCE k instead of J NCE is so that we can incorporate temporal correlation between t and t − k. Another motivation is that we can increase the model capacity of the discriminator p(z t |x t ) by incorporating the independent dynamics modelp(z t |·).

Model Definitions
This section discusses how we define the discriminator components: p(z t |x t ) andp(z t |z t−1 , a t−1 ). Ref. [34] has empirically shown that the inductive bias from model architectures is a significant factor for contrastive learning. As experimentally recommended in the literature, we define p(z t |x t ) as an exponentiated bilinear similarity function parameterized with W z|x : where e t := f CNN (x t ) and f CNN (·) denotes feature extraction by a CNN unit. With this definition, J NCE k is simply a softmax cross-entropy objective with logits z W e. Contrary to the previous contrastive learning literature [22,34], the definition of newly introducedp(z t |z t−1 , a t−1 ) is required. Here, we propose to apply linear modeling to define the model deterministically as: δ is the Dirac delta function, and W z,a are linear parameters.
This linear modeling ofp(z t |·) successfully regularizes J NCE k and contributes to construct smooth latent space. We can alternatively definep(z z |·) := p(z z |·), where p(z z |·) is generally defined as an expressive model aiming at precise prediction. However, the high model capacity allows to embed temporally consecutive samples too distant from each other to sufficiently optimize J NCE k , thus yielding unsmooth latent space. Figure 3 illustrates the architecture to compute J NCE k . We describe the two paramount components which characterize our proposed contrastive learning scheme as follows:

Instantiation with RSSM
(i) Independent linear forward dynamics: As previously proposed in Sec. 4.3, we employ a simple linear forward dynamicsp, which is used only for contrastive learning. During the policy optimization phase, the expressive model with GRU is alternatively utilized to make the most out of its long-term prediction accuracy.
(ii) Data augmentation: We append two independent image preprocessors which process two sets of input images (i.e., x ≤t and x t+1:t+K ). Considering the empirical success of the previous literature [23,25,28,30,29], we adopt the random crop of images. In our implementation, the original   (64, 64). The origin of the crop rectangle is determined at each preprocessor randomly and indenpendently. This makes it difficult for the contrastive learner to discriminate correct positive pairs, encouraging only informative features for control to be extracted.
We propose a decoder-free variant of Dreamer, which we call Dreamer with InfoMax and without generative decoder (Dreaming). Dreaming trains a policy as almost same way with the original Dreamer. The only difference between the methods is that we alternatively use the contrastive learning scheme introduced in the previous section to train RSSM. We implement Dreaming in TensorFlow [35] by modifying the official source code of Dreamer 1 . We keep all hyperparameters and experimental conditions similar to the original ones. A newly introduced hyperparameter K in Eq. (8) (overshooting distance) is set to be K = 3 based on the ablation study in Appx. D.

Comparison to State-of-the-art Methods
The main objective of this experiment is to demonstrate that Dreaming has advantages over the baseline method Dreamer [16] on difficult 5 manipulation tasks exhibited in Fig. 1, in which Dreamer suffers from object vanishing. The specifications of the two original tasks, UR5-reach and Connectorinsert, are described in Appx. B. For the difficult 5 tasks, we also compare the performance with the latest cutting edge MFRL methods, which are CURL [25], DrQ [29] and RAD [28]. In addition, another variety of 10 DMC tasks are evaluated, which are categorized into three classes namely; manipulation, pole-swingup, and locomotion. For the additional 10 tasks, only CURL is selected as an MFRL representative. Table 1 summarizes the training results benchmarked at certain environment steps. Their training curves are shown in Appx. C. (A) We put much focus on these difficult tasks and it can be seen that Dreaming consistently achieves better performance than Dreamer. Hence, this indicates that the decoder-free nature of the proposed method successfully surmounts the object vanishing problem. In addition, Dreaming achieves outperforming performance than the leading MFRL methods. (B)  . On these tasks, the cameras always track the center of locomotive robots, and this causes the key control information (i.e., velocity) to be extracted from the background texture. We suppose that this robot-centric nature makes it difficult for the contrastive learner to extract such information because only robots' attitudes provide enough information to discriminate different samples. Figure 4 shows video prediction by Dreaming, in which principal features for control (e.g., positions and orientations) are successfully reconstructed from the embeddings learned without likelihood objective. However on Cheeta-run, another kind of object vanishing arises; the checkered floor pattern, which is required to extract the velocity information, is vanished.

Ablation Study
This experiment is conducted to analyze how the major components of the proposed representation learning, introduced in Sec. 4.4, contribute to the overall performance. For this purpose, some variants of the proposed method have been prepared: (i) the effect of independent linear dynamics is demonstrated with a variant that has shared dynamicsp(z t |z t−1 , a t−1 ) := p(z t |z t−1 , a t−1 ), (ii) the effect of data augmentation is demonstrated by removing the image preprocessors shown in Fig. 3. We also prepare another data augmentation called color jittering [23,29,28], for reference. Only three tasks, which are Cup-catch, Reacher-hard, and Finger-turn-hard, are taken into this experiment. Tables 2 and 3 summarize the results from the performed ablation study, which reveals that both of the proposed components are essential to achieve state-of-the-art results.

Conclusion
In the present paper, we proposed Dreaming, a decoder-free extension of the state-of-the-art MBRL method from pixels, Dreamer. A likelihood-free contrastive objective was derived by reformulating the original evidence lower bound of Dreamer. We incorporated the two indispensable components below to the contrastive learning: (i) independent and linear forward dynamics, (ii) the random crop data augmentation. By making the most of the decoder-free nature and the two components, Dreaming was able to outperform the baseline methods on difficult tasks especially where Dreamer suffers from object vanishing. We believe that the proposed representation learning scheme is applicable to other MBRL methods [14,17,15,31].
An disadvantage we observed in the experiments was that Dreaming degraded the training performance on planar locomotion tasks (e.g., Walker-walk), where the contrastive learner has to focus on not only robots but also the background texture. This weak point should be resolved in future work as it may affect industrial manipulation tasks where first-person-view from robots dynamically changes. Another future research direction is to incorporate the uncertainty-aware concepts proposed in recent MBRL studies [1,2,19,30]. Although we have achieved state-of-the-art results on some difficult tasks, we have often observed overfitted behaviors during the early training phase. We believe that this model-bias problem [36] can be successfully solved by the above state-of-the-art strategy.

A Derivation
In this section, we clarify that J NCE k is a lower bound of I(x t , z t−k ). J NCE k can be rewriten as: where f (x t , z t−k ) includes deterministic multi-step prediction withp(z t |z t−k , a <t ) and computation of the bilinear similarity by Eq. (9). For ease of notation, actions a <t in the conditioning set are omitted from f (·). As already shown in [22], the optimal value of f (·) is given by; By applying Bayes' theorem f (x t , z t−k ) ∝ p(z t−k |x t )/p(z t ) and inserting this to Eq. (11), we get; By marginalizing Eq. (13) with respect to the data distribution, we finally obtain; Note that setting k = 0 derives E[J NCE ] ≤ I(x t ; z t ). Connectorinsert (right) is to insert a millimeter-sized connector gripped by a robot to a socket. This tasks is originally introduced in [37]. Since the gap between the connector and socket is very tight, pixelwise precise control is required. In the both tasks, the goal positions are initialized at random. Reward functions are defined with tolerance() function [20] as: tolerance(x, bounds=(0, 0.04), margin=0.5) for UR5-reach and tolerance(x, bounds=(0, 0.5), margin=1.) for Connector-insert, respectively. x denotes distance to goal in meters (UR5-reach) and millimeters (Connector-insert).