Augmenting Replay in World Models for Continual Reinforcement Learning

A Recurrent State-Space Model (RSSM) [6], predicts environment dynamics, Figure 1. It maintains a deterministic hidden state $h_{t}$ and models a stochastic representation of the next state given current and previous observations $x_{1:t}$ and actions $a_{1:t}$ Dynamics are modelled with a Gated Recurrent Unit (GRU) network, predicting the deterministic state $h_{t+1}=\text{GRU}(h_{t},z_{t},a_{t})$ and the stochastic state $\hat{z}_{t+1}=f(h_{t+1})$. The stochastic state $z_{t}$ is either inferred by a variational autoencoder $z_{t}\sim q_{\theta}(z_{t}\mid h_{t},x_{t})$ or predicted by the dynamics model $\hat{z}_{t}\sim p_{\theta}(\hat{z}_{t}\mid h_{t})$ for open-loop prediction in dreaming, where stochastic state posteriors $z_{t}$ are unavailable.

We used a standard GRU with Tanh activation for the recurrent component. The stochastic state $z_{t}$ comprises 32 discrete stochastic units, each with 32 categorical classes, following the architecture of DreamerV3.

The world model state at timestep $t$ is the concatenation of $h_{t}$ and $z_{t}$, forming a Markovian representation of the environment state. The world model is trained to reconstruct input images and rewards. KL balancing [7] helps to model the transition between states.

The actor and critic (S.LABEL:app:network_arch) are MLPs that map the state of the model to actions and value estimates. They are trained entirely on trajectories generated stochastically by the world model, referred to as ‘dreaming’. The actor and critic are trained on-policy using REINFORCE [23]. Inevitable changes in the model state space are not an issue since imagined trajectories are cheap to generate and do not require interaction with the actual environment, allowing this process to be run to convergence.

The augmented replay buffer, Figure 2, comprises a short-term FIFO buffer $\mathcal{D}_{1}$ (as used in DreamerV3) and a long-term global distribution matching buffer $\mathcal{D}_{2}$. They are equally sized and used in parallel. Data from both are uniformly sampled for each training minibatch (S.Algorithm LABEL:alg:combined_buffers). A key objective is to minimise memory requirements and hence the size of the replay buffer. While Dreamer maintained a single buffer containing the last 1 million (1,000,000) observations, we empirically chose a size of $2^{18}\approx$ 262,000 observations for each buffer, resulting in an augmented buffer that is significantly smaller (about $2\times$), without noticeably affecting performance. We also introduced spliced rollouts, which are a simple and sometimes necessary alternative to storing entire episodes, enabling smaller buffer sizes.

In tasks without shared structure, there are significant differences between environments in terms of task dynamics, visual differences, and magnitude of rewards, which poses a significant challenge to CL, especially without task id’s. Exploring new environments is difficult, as policies that have already been trained on previous tasks may lack the randomness required to adequately explore the state-space of a new task. We addressed this challenge using fixed-entropy regularisation (as in DreamerV3) for training the world model. In addition, we used predetermined reward scales in individual environments based on single-task baselines, during actor-critic training. This approach enabled mitigation of the exploration challenge without the addition of an exploration-orientated learning system such as Plan2Explore [22].

We evaluated WMAR in a set of challenging OpenAI Procgen and Atari environments [3, 1]. They are commonly used RL benchmarks, can be extremely similar for the case of shared structure (Procgen) or dissimilar where there is no shared structure (Atari), and they were computationally feasible on the available hardware. We define shared structure as tasks that have similar environment state and observations, action dynamics, and rewards. We applied visual perturbations to create differences between tasks. To evaluate performance on tasks without shared structure, we selected 4 Atari tasks.

To measure the efficacy of the augmented replay buffer (i.e., adding a long-term distribution matching buffer with spliced rollouts), we compared WMAR and DreamerV3, with equal memory allowance for replay buffers. For the most direct comparison, we implemented our own version of DreamerV3, referred to as $DV3^{\prime}$, and augmented the replay buffer to create WMAR. To validate DV3’, we compared to an open source implementation of DreamerV3 https://github.com/danijar/dreamerv3, see S.LABEL:app:validation.

We also ran single-task baselines: WMAR and a random agent. The single-tasks baselines were used for evaluation metrics.

We evaluated agents using task reward and following Kessler et al. [12], extended the evaluation to forgetting and forward transfer. Forgetting informs us about backward transfer and stability, and forward transfer informs us about plasticity.

We evaluated performance by normalising episodic reward to two single-task baselines: WMAR and a random agent. The normalised comparison allowed us to measure the relative performance, providing a natural evaluation of forgetting and forward transfer for all environments in each CL suite.

All WMAR experiments are trained on one Nvidia A40 40GB GPU with all single-task benchmarks running within 0.25 days and continual learning benchmarks running within 1 day of wall time, making these experiments widely available and reproducible across research labs. For more details, see S.LABEL:app:extime.

We chose a subset of Atari environments where the agent could achieve reasonable performance with less training and followed [15] in using sticky actions. The environments were presented with arbitrary ordering, as testing each permutation would have been prohibitively expensive.

Normalised performance is plotted in Figure 3, and single-task results used for reward scaling are given in S.LABEL:app:baseline. The predetermined reward scales are shown in S.Table LABEL:tab:atari_rew_scales. WMAR was capable of continual RL, whereas it is clear that $DV3^{\prime}$ was not; it only performed well when it was actively training in a given task.

Normalised performance is shown in Figure 4, and single-task results used for reward scaling are given in S.LABEL:app:baseline. The terminology for CoinRun environments is: NB = no background, RT = restricted themes, and MA = monochrome assets.

Although CoinRun environments have identical mechanics, they share both subtle and substantial visual differences, posing significant challenges. However, WMAR could learn continuously and displayed forward and backward transfer of skills between environments. $DV3^{\prime}$ showed similar characteristics, but performance dropped on the first task, ‘CoinRun’, after training on it (forgetting), and on the last task before training on it (lacking forward transfer). They had approximately equal peak performance on each task; however, WMAR was noticeably more consistent than $DV3^{\prime}$. Overall, WMAR showed more consistent performance across tasks with fewer fluctuations in performance during training, which may be a desirable property in practical applications.

This study provides a strong justification for future work; however, the experiments were limited in scope. They could be expanded (e.g., as in [5, 12]) by increasing the length of task sequences, increasing the number of environment steps, testing on more environments, and comparing to additional reference models.

Testing all task permutations is prohibitively expensive and therefore, we used a set task order. However, task ordering can result in significant differences in CL results (e.g., Appendix G in [20]). Future work could randomise the order (as a proxy for testing all permutations) over a larger number of seeds.

Another area for future work is to improve WMAR by combining it with existing techniques such as behaviour cloning in CLEAR. Such an adaptation could counter shifts in the latent distribution as the world model trains and where the actor is frozen.