Fractional Transfer Learning for
Deep Model-Based Reinforcement Learning

We can formalize an RL problem as a Markov Decision Process (MDP) (Bellman, 1957), which consists of four main elements:

• •

  a set of states $S$

• •

  a set of actions $A$

• •

  a transition function $P(s,a,s^{\prime})$ = Pr ($s_{t+1}=s^{\prime}|s_{t}=s,$ $a_{t}=a$)

• •

  a reward function $R(s,a,s^{\prime})$

Here $P(s,a,s^{\prime})$ denotes the probability that taking action $a\in A$ in state $s\in S$ at time step $t$ results in a transition to state $s^{\prime}\in S$. When a state transition takes place, the reward function $R(s,a,s^{\prime})$ yields an immediate reward corresponding to this transition. The objective in an MDP is to find an optimal policy $\pi^{*}(s_{t})=a_{t}$ that maximizes the expected cumulative reward. The expected cumulative reward is defined as:

Here $\gamma$ represents the discount factor (where $\gamma\in$ [0,1)), and $T$ is a random variable representing the number of time steps to account for in computing the expected cumulative reward.

The number of states $|S|$ and actions $|A|$ may be very large or infinite, in which case neural networks have typically been used as function approximators for the past decades (Tesauro, 1995). In more recent years, DL architectures such as Convolutional Neural Networks (CNNs) (LeCun et al., 2015) and Recurrent Neural Networks (RNNs) (Medsker & Jain, 1999) also sparked the interest of the RL community, which resulted in Deep RL (DRL) research. DRL became particularly popular since the introduction of Deep Q-Networks (DQN) (Mnih et al., 2015), where CNNs were used to process high-dimensional raw image data. This resulted in agents achieving a level comparable to professional human players in 49 Atari games, solely based on a reward signal and image observations. Significant advancements that followed included the iterations of AlphaGo (Silver et al., 2016), AlphaZero (Silver et al., 2017), and eventually MuZero (Schrittwieser et al., 2020). The latter accomplished state-of-the-art performances in Go, Chess, Shogi, and the Atari benchmarks by making use of MBRL. DRL has also been shown to be able to master highly complex modern games, such as Starcraft II (Vinyals et al., 2019) and Dota 2 (Berner et al., 2019).

Whereas model-free methods can directly learn from high-dimensional pixel images, MBRL faces the difficulty that a transition model is hard to train from such high-dimensional data. MBRL was shown to be efficient and successful for small sets of low-dimensional data, such as PILCO (Deisenroth et al., 2013), where the authors used Gaussian processes (GP) to learn a model of the environment from which trajectories could be sampled. However, methods such as GP don’t scale to high-dimensional observations. One work showed that by first using an autoencoder (AE) (Ballard, 1987) to learn a lower-dimensional representation of image data, a controller could be trained to balance a pendulum using the feature vectors provided by the bottle-neck hidden layer of the AE (Wahlström et al., 2015). This usage of compressed latent spaces allowed much more sample efficient RL methods to be developed (Finn et al., 2016; Watter et al., 2015).

More recently, World Models (Ha & Schmidhuber, 2018) were introduced that allowed MBRL agents to learn policies for visual environments without having to interact with the environment. The authors first trained a Variational AE (VAE) (Rezende et al., 2014; Kingma & Welling, 2014) to learn a compressed representation of randomly gathered image data. Opposed to learning a function to compress the inputs, as is done with vanilla AEs, VAEs learn the parameters of a probability distribution representing the data. Using the VAE produced latent space, an RNN was trained that takes as input a latent state $z_{t}$, an action $a_{t}$, and a hidden state $h_{t}$ to predict a probability distribution of the next latent state $z_{t+1}$. The RNN is combined with a Mixture Density Network (MDN), such that a temperature parameter can be introduced that allows controlling the model’s uncertainty when sampling from the distribution. After training, the MDN-RNN model allowed ’imagined’ trajectories to be sampled that represent the actual environment. Additionally, the authors could assume access to the underlying reward functions of the used environments. A linear controller policy could therefore be learned without interacting with the actual environment, after which the policy was able to be transferred to the actual environment. The authors of SimPLe extended this approach, where they showed that a World Model can be used as a simulator to train model-free policies, resulting in impressive sample efficiency (Kaiser et al., 2019).

Rather than learning a World Model and policy in separate stages, PlaNet (Hafner et al., 2019b) trains these processes jointly and performs latent online planning for visual continuous control tasks. Moreover, instead of assuming the reward function of the environment, PlaNet learns a model of the reward as a function of the latent states. The latter allows planning to be performed completely in imagination. For planning, the authors combine model-predictive control with the cross entropy method. The latter is a population-based optimizer that iteratively evaluates a set of candidate action sequences under the model that are sampled from a belief distribution, which is repeatedly re-fit with the top $K$ rewarding action sequences. As transition model architecture, the authors introduce the Recurrent State Space Model (RSSM). This model combines a deterministic and stochastic model into one architecture. Purely deterministic models, such as a vanilla RNN, cannot capture multiple futures, and tend to exploit model inaccuracies. Purely stochastic models have difficulties remembering information over multiple time steps, given the randomness in transitions. The authors split the latent state in stochastic and deterministic parts, allowing robust prediction of multiple futures. This work provided the foundation for the state-of-the-art MBRL continuous control algorithm Dreamer (Hafner et al., 2019a), as well as its variant that is adapted for discrete problems, DreamerV2 (Hafner et al., 2020). Dreamer will be the main focus of this work, and a detailed explanation of the core components and processes can be found in Appendix A.

This section introduces the concept of TL in the RL domain. We define TL and detail on how it is often used for RL algorithms (Section 2.3.1). Additionally, we cover prior research that have applied TL to MBRL, with a main focus of approaches using World Models (Section 2.3.2).

As was shown in PlaNet (Hafner et al., 2019b), a single agent using the RSSM World Model is able to learn multiple tasks from different domains simultaneously if the dynamics across the domains are similar, and visual observations are used to be able to infer which domain is being faced. In this work we investigate whether these simultaneous multi-task agents can serve as a new type of multi-source TL for RL. In order to train agents on multiple tasks from different domains simultaneously, some adjustments need to be made to Dreamer, as was done in the PlaNet paper. First, when dealing with multiple tasks simultaneously, the action dimension $|A|$ of each domain may be different. Components like the action model in Dreamer require $|A|$ to be fixed in order to provide compatible outputs for each domain. To overcome this issue, we first compute which task has the largest action space dimensions, which is then used as the target dimension for padding the action spaces of all other tasks with unused elements:

Where $D\in\textbf{D}$ is domain $D$ from the set of domains D that the agent will be trained in. This means the action space for each domain will be equivalent, but for tasks with originally smaller action spaces there are elements that have no effect on the environment and can therefore be ignored by the agent. Next, in order to acquire a balanced amount of gathered experience in the replay dataset, we apply the following modifications. Instead of randomly collecting 5 episodes for initialization, we randomly collect 1 episode of each task. Additionally, instead of collecting 1 episode of a single task after 100 training steps, we collect 1 episode for each of the training tasks.

As discussed in Section 2.3, when transferring representations of neural networks the common approach is to re-use all weights of each layer. The final layer often needs to be randomly initialized for learning new tasks, as it otherwise needs to unlearn the previously fit weights. The other layers are either frozen, or retrained together with the final layer. However, the drawback of initializing a layer randomly is that all previously obtained knowledge contained in that layer is lost. Therefore, we present a novel alternative TL approach: FTL. The idea is to transfer a fraction of the weights of a source layer to a target layer. This can be done adding a fraction of the source layer’s weights to the randomly initialized weights of a target layer. Rather than discarding all knowledge by random initialization, this approach allows the transfer of portions of potentially useful knowledge. Additionally, we can avoid overfitting scenarios with this approach as we don’t necessarily fully transfer the parameters.

Formally, we can simply define FTL as:

Where $W_{T}$ is the randomly initialized weights of a target layer, $W_{S}$ the trained weights of a source layer, and $\omega$ the fraction parameter that quantifies what proportion of knowledge is transferred, where $\omega\in[0,1]$. That is, for $\omega=0$ we have no transfer (i.e. random initialization), and for $\omega=1$ we have a full transfer of the weights.

The experimental goal of this study is to transfer knowledge from agents that were simultaneously trained on multiple tasks to novel tasks. We are interested in observing whether this provides for a positive transfer, and additionally whether FTL provides for even better performances. However, as Dreamer consists of several different components, we need to identify what type of TL is most suitable for transferring each component’s learned representation. In Figure 1 a schematic illustration of all components of Dreamer and what type of transfer is applied to them in this work can be observed, as will be discussed below. The reward, value, and action models are each implemented as three fully-connected layers, the encoder and decoder are four-layer CNNs, and the dynamics model is the RSSM as described in Section 2.2. All of the following described parameters that are transferred are kept trainable afterwards.

First, we identify the components of Dreamer that require random initialization, which are the components that involve processing actions. This is due to the fact that when transferring to a new domain, the action dimensions are likely incompatible with the action dimensions from the multi-source agents. Moreover, even if the action dimensions are equivalent, when the source and target domains are different, the learnt behaviours won’t apply to the target domain and have to be unlearned when fully transferred. Therefore, we chose to randomly initialize (i.e. $\omega=0$) the last layer of the action model, and the weights that connect the input actions to the forward dynamics model.

Next, we identify the components that are suitable for FTL, being the value model and reward model. For these neural networks we are not dealing with dimension mismatches, meaning random initialization is not a necessity. Additionally, full transfers are not applicable, as reward and value functions are generally different between the source and target domains, which would mean the weights of the source task would have to be unlearned. However, though these functions may not be equivalent across domains, the learned weights may still contain useful information that can be transferred with FTL. For these reasons we perform FTL for the last layer of the reward model and the last layer of the value model.

For the action model, value model, and reward model we only perform FTL and random initialization for the last layer. For the previous layers of these models we apply full transfer, as is generally done in TL for supervised learning, where only the last layer is randomly initialized. As the previous layers solely learn feature extractions of the latent state inputs, they can be re-used across domains. Preliminary empirical results showed that the best performances were obtained when only treating the last layer differently.

Finally, we identify which components of Dreamer are suitable for full transfers, which are the forward dynamics model and the VAE. The forward dynamics model can likely be fully transferred across domains if, for instance, the physical laws across the domains are equivalent, meaning this type of knowledge can be re-used (Taylor & Stone, 2009). The domains that will be used in this study share the same physics engine (Section 4), meaning we decided to completely transfer the dynamics model parameters. We can also fully transfer the parameters of the representation model, as the generality of convolutional features allow reconstruction of observations for visually similar domains, and can therefore quickly adapt to novel domains (Chen et al., 2021).