Revisiting “Recurrent World Models Facilitate Policy Evolution”

Abstract

This paper contributes a detailed analysis of the architecture of Ha and Schmidhuber [5]. The original paper proposes an architecture comprising 3 main components: a “visual” module, a “memory” module; and a controller. As a whole, such architecture performed well in challenging domains. We investigate how each of the aforementioned components contributes individually to the final performance of the system. Our results shed additional light on the role of the different components in the overall behavior of the agent, and illustrate how the different design options affect the behavior of the resulting agent.

This work was supported by national funds through Fundação para a Ciência e a Tecnologia (FCT) with ref. UIDB/50021/2020. BE acknowledges a research grant from Fundação Calouste Gulbenkian under program “Novos Talentos em IA.”

Appendices

AFull Comparative Results

For completeness, we provide in this appendix some additional results not reported in the main body of the paper.

1.1 A.1Ablation Study Additional Results

We include, in Table 6, a comparison between the average score obtained when using the full image and a cropped image. The results complements those portrayed in Table 2, but suggest essentially the same conclusions.

Table 6. Average score on both the game’s cropped and full image

Table 7. Average score obtained with $V_\mathrm{mean}$ trained with images from different sampling methods on the game’s full image.

Table 8. Average score and steps to reach 900-score, training the VAE and MDN-RNN components with images obtained using different sampling methods.

Fig. 7.

VAE architectural details from the original World Models implementation [5]

1.2 A.2Improved Sample Policy

The VAE and MDN-RNN components are trained with a batch of images from the game sampled with a random driving policy. However, it is also important to understand the impact that the policy used to sample the batch of images used to train these components has on the performance of the system as a whole. We thus considered training these components with a mix of images sampled with both a random policy and an expert policy. Results obtained with such a mix sample are denoted as improved sample.

Table 7 shows the performance obtained when the visual component is trained with images from a mixed policy, in contrast with the performance obtained with images from a random policy. We can observe an improvement both in the cropped and the full image cases.

In Table 8 reports the results obtained when we include a memory component. We report both the average score and the time step at which the agent attained a score of 900. Using a better sampling method for training the “visual” and “memory” components, the system as a whole tends to achieve better results—mostly when using cropped images—and reach a high score faster.

BModel details

In this section we provide the details of the architectures used in the different experiments. The overall structure of the neural networks used closely follow those of the original HS paper [5], and we refer to that work for further details.

1.1 B.1 VAE

We start by presenting the details for the VAE. The encoder and decoder are represented in Fig. 7. The encoder comprises 4 convolutional layers where, in each layer, the filter size is always 4, and the stride is 2. The dimension of the latent space is $N_z=32$. The decoder, in turn, comprises 4 deconvolution layers that reconstruct a $64\times 64\times 3$ image $\hat{\boldsymbol{x}}$ from a code $\boldsymbol{z}$. The VAE is trained using gradient descent with a dataset of previously gathered environment frames.

1.2 B.2 MDN-RNN

The memory module consists of a recurrent neural network (RNN) followed by a mixture density network (MDN) like represented on Fig. 4. The RNN is a single LSTM [7] layer with a hidden size of 256 and the MDN is composed by a single linear fully-connected layer with 5 gaussians per latent unit of $\boldsymbol{z}$. The MDN-RNN is trained with gradient descent, by using the previous learned VAE and trying to sequentially reconstruct the sequences from the same previously garthered environment dataset.

1.3 B.3 Controller

The original controller presented by Ha et al. [5] consists of a single fully-connected layer with an custom environment specific activation function defined in Table. 9. In this work, we also analyzed different architecture choices like removing the squashing activation function and adding a hidden layer with 24 hidden units and with Tanh activation function. The controller is trained with CMA-ES [6] by interacting with the environment.

Table 9. Custom activation function from Ha et al. [5] implementation

1.4 B.4 Hyperparameters

(See Table 10)

Table 10. Hyperparameters used for training the World Models agent