# Boosting Deep Reinforcement Learning Agents with Generative Data Augmentation
## 1 Introduction  
Recently, deep learning models have demonstrated a significant performance boost when data augmentation techniques are employed, mainly in tasks where it is difficult to obtain high-quality actual data [1; 2]. In such cases, synthetic samples are produced out of real ones, in order to tackle dataset imbalance and enlarge small datasets, as the efficiency of deep learning algorithms is highly dependent on the data used during training. Under this scope, several methodologies have been proposed, mainly in the computer vision domain, ranging from geometric transformations [3; 4] to more advanced deep learning approaches [5; 6; 7]. Tasks such as image classification and object detection are shown to benefit significantly from data augmentation, highlighting the importance of diverse, high-quality training samples. In the context of reinforcement learning (RL), the generalization and convergence speed of the agent is one of the main subjects of research, as many approaches struggle to efficiently utilize the available data [8; 9], especially in environments with sparse rewards. Thus, several attempts have been made to augment the datasets used for the agent's training, achieving quite satisfactory results and indicating that data augmentation is a promising direction in RL [10; 11; 12; 13]. Specifically, in RL environments with visual observations (where the state representation is an image), most existing approaches aim to augment data by processing the original images with traditional techniques such as cropping, flipping, etc. [12; 14]. Although, these methods increase the model's generalization ability, they lack diversity in the sense that the new data arise directly from existing samples. In this work, a framework is presented for creating novel synthetic samples by learning the distribution of the original images (states) rather than processing existing observations. The main idea is to generate observations as similar as possible to states the agent could encounter in the future and use them at earlier stages of the agent's training. Depending on the examined environment, states that provide useful information (basically in terms of obtained rewards) may be observed after many timesteps, delaying the progress of training. On top of that, the learned policy determines, to a high degree, the observed states. This interdependence between the policy and the training samples may lead to local optima and prevent the agent from improving its policy. In order to overcome this issue, a generative process is proposed to enhance the diversity of the samples and speed up training by mixing synthetic and real samples in each batch. Another environment-specific issue that may arise in RL tasks is the imbalance of the training set (with respect to the rewards of the collected samples). Similarly to class-imbalanced supervised tasks, in environments with sparse rewards the agent is trained on a dataset in which the majority of samples are neutral (i.e., the immediate reward is zero), making the distinction among the qualities of different actions quite complex. Under this scope, generative augmentation can be used to enrich the original set (consisting of samples collected by the agent during its interaction with the environment) with samples of a specific category (e.g., states with high reward). Based on the technique introduced in this paper, samples can be produced independently of the type of observed data and augment the dataset towards a certain direction, limiting the effect of class imbalance in the agent's performance. The proposed methodology is split into two phases. At first, a diffusion model is used to generate images that represent the current state and the next state of a sample. Subsequently, a custom model for predicting the agent's action that was executed in between is employed, while the reward is determined based on the data that were used for the training of the diffusion model.