# Boosting Deep Reinforcement Learning Agents with Generative Data Augmentation
## 4 Augmentation Framework
### Data Generation  
The process of generating complete synthetic samples of the form \((s,a,r,s^{\prime})\) is depicted in Figure 2. Two DDPMs are trained on samples of two consecutive game states in order to produce a state \(s\) and the next state \(s^{\prime}\). The goal is to generate samples with either very high or very low rewards since the original dataset created by the agent's observations is very imbalanced in the sense that the immediate reward in most cases is zero. The training dataset for the first model consists of pairs \((s,s^{\prime})\) with a high reward, while the second one contains pairs with a low reward. As each model is trained on a specific type of data, we can assume that the states generated by the first model correspond to high rewards while the states generated by the second model correspond to low rewards. In this way, the reward for each synthetic sample is considered as known by the time of its generation.  
Figure 2: Synthetic data generation.  
Since the executed action is not taken into consideration during DDPM training, each generated pair \((s,s^{\prime})\) could be the result of a different action. In order to produce functional samples that can be used to train the agent, a separate model has been employed for action prediction given two consecutive states. Thus, for each pair \((s,s^{\prime})\) of synthetic states, the action is predicted by this model and the reward is immediately set to either \(1\) or \(-1\) depending on the DDPM used for the generation, resulting in a complete training sample. The described process is presented step by step in Algorithm 1.  
```
1 FunctionGenerateSyntheticData():
2\(action\_predictor\_buffer,d\_dpm\_buffer\_high,d\_dpm\_buffer\_low\leftarrow[],[],[]\)while\(action\_predictor\_buffer.size()<max\_size\)do
3\(a\gets agent.select\_action(s)\)\(s^{\prime},r\gets env.execute(a)\)\(action\_predictor\_buffer.append(s,a,r,s^{\prime})\)if\(r>0\)then
4\(d\_dpm\_buffer\_high.append(s,a,r,s^{\prime})\)if\(r<0\)then
5\(d\_dpm\_buffer\_low.append(s,a,r,s^{\prime})\)\(s\gets s^{\prime}\)\(action\_predictor.train()\)\(d\_dpm\_high.train()\)\(d\_dpm\_low.train()\)\(\{s_{b},s^{\prime}_{b}\}\gets d\_dpm\_high.generate\_states()\)\(\{s_{t},s^{\prime}_{b}\}\gets d\_dpm\_low.generate\_states()\)\(\{a_{pred}\}\gets action\_predictor.predict(\{s_{h},s^{\prime}_{h}\},\{s_{l},s^{\prime}_{l}\})\)\(synthetic\_data.append(\{s_{h},a_{pred},r,s^{\prime}_{h}\})\)\(synthetic\_data.append(\{s_{l},a_{pred},r,s^{\prime}_{l}\})\)return\(synthetic\_data\)
```  
**Algorithm 1**Synthetic data generation  
Figure 3 shows samples generated with the methodology described above for the ALE boxing environment. Specifically, frames \(h_{0}\)- \(h_{4}\) are generated with the diffusion model trained on states that lead to high rewards, while frames \(l_{0}\)-\(l_{4}\) are produced by the second DDPM and represent two consecutive states with a low immediate reward. The synthetic images are then fed to the action-predicting model in order to be paired with the most suitable action.  
It is evident that the generated images are quite similar to the real ones in terms of quality. The generative models successfully produce consecutive frames while maintaining the more general features of the image. As depicted, the two models are able to generate  
Figure 3: Synthetically generated samples for the ALE boxing game.  
meaningful states in the sense that the transition between two consecutive frames is smooth, but also capture the most crucial patterns that lead to either positive or negative rewards. In this manner, the produced states are functional and there is no risk of misguiding the training procedure by augmenting the dataset with irrelevant images.