Neighborhood Mixup Experience Replay: Local Convex Interpolation for Improved Sample Efficiency in Continuous Control Tasks
Abstract
Experience replay plays a crucial role in improving the sample efficiency of deep reinforcement learning agents.
Recent advances in experience replay propose the use of Mixup [35] to further improve sample efficiency via synthetic sample generation.
We build upon this idea with Neighborhood Mixup Experience
Replay (NMER), a modular replay buffer that interpolates transitions with their closest neighbors in normalized state-action space.
NMER preserves a locally linear approximation of the transition manifold by only performing Mixup between transitions with similar state-action features.
Under NMER, a given transition’s set of state-action neighbors is dynamic and episode agnostic, in turn encouraging greater policy generalizability via cross-episode interpolation.
We combine our approach with recent off-policy deep reinforcement learning algorithms and evaluate on several continuous control environments.
We observe that NMER improves sample efficiency by an average 87% (TD3) and 29% (SAC) over baseline replay buffers, enabling agents to effectively recombine previous experiences and learn from limited data.
Introduction
Learning robust and effective behavior from a limited set of examples is a hallmark of human cognition [6].
The sample efficiency of our brain and neuronal circuits allows us to quickly learn new skills, perform new tasks, and generalize the experience we observe to a variety of different environments and settings.
In many problem domains, human sample efficiency far outperforms that of deep reinforcement learning algorithms [14].
Narrowing this gap is a critical milestone towards effectively replicating intelligence in reinforcement learning settings.
Model-free (MF), off-policy deep reinforcement learning (DRL) algorithms provide significant sample efficiency gains relative to their on-policy counterparts [8, 2].
Improved sample efficiency is due in part to experience replay [18] techniques, which enable agents to continuously learn from past experiences.
The trial-and-error nature of reinforcement learning nonetheless still necessitates the need to collect large volumes of training data [32, 33, 17].
While lower sample efficiency may be acceptable for learning in simulation, it may significantly hinder an agent’s progress in many real-world applications where samples are expensive to generate [32].
Model-based (MB)
DRL agents achieve improved sample efficiency by learning a model of the environment [30, 9, 11] that can be used for offline planning and policy refinement [4, 24, 25].
For reinforcement learning tasks with noisy and high-dimensional state and action spaces, however, an agent’s learned environment models may suffer from estimation error and bias where little data is available [17, 22].
Combined with model capacity limitations, this can result in model-based agents converging to suboptimal policies.
We aim to combine the benefits of learning on true environment interactions in MF-DRL with the sample efficiency benefits of MB-DRL.
To this end, we propose Neighborhood Mixup Experience
Replay (NMER), a modular replay buffer that improves the sample efficiency of MF, off-policy DRL agents by training on experiences sampled from convex, locally linear interpolated transitions of the replay buffer.
At a high level, NMER interpolates proximal pairs of transitions in the transition space of the replay buffer using Mixup [35], a convex and stochastic linear interpolation technique.
Despite the computational and analytical simplicity of Mixup, which enables for runtime-efficient experience replay, the use of Mixup in experience replay can enable better generalization and policy convergence in MF, off-policy DRL agents through implicit regularization and expansion of the training support of the agent’s function approximators.
We empirically observe these beneficial effects of Mixup in our continuous control experiments.
NMER can be applied to any continuous control reinforcement learning agent leveraging a replay buffer.
As a motivating example, consider a robotic humanoid learning to walk using MF, off-policy
DRL, given a limited set of experiences consisting of odometry and actuator sensor measurements.
With standard experience replay [18] approaches, the finite size of the agent’s experiences can hinder the agent from learning robust policies, perhaps due to the agent not experiencing a crucial subset of the transition space.
With NMER, however, interpolated experiences can provide this DRL agent with crucial training samples in regions of the transition space they previously did not experience, thus improving the robustness and performance of the policies the agent learns.
Our contributions are summarized as follows:
Neighborhood Mixup Experience Replay (NMER) :
A replay buffer that improves sample efficiency of off-policy DRL agents by training these agents on locally linear interpolated combinations of neighboring transitions.
Local Mixup : A generalization of NMER, this algorithm considers the application of Mixup between local points in any feature or vector space by only applying Mixup to nearest neighbors measured using a given distance metric.
In highly nonlinear features spaces, local Mixup enables for learner generalization and regularization without interpolating too far across highly different dynamic regions of the transition space.
Improved sample efficiency in continuous control : Our evaluation study demonstrates that NMER substantially improves sample efficiency of MF, off-policy DRL algorithms across several continuous control environments.
Related work
Experience replay, data augmentation, and interpolation approaches have been applied to RL and other machine learning domains.
NMER builds off of these techniques to improve sample efficiency.
Experience replay Prioritized Experience Replay (PER) [23] samples an agent’s experiences from a replay buffer according to the “learnability” or “surprise” that each sample induces in the agent in its current parameterization.
PER uses absolute TD-error of a sample [23] as a heuristic measure of “surprise”.
In the stochastic prioritization variant of PER, transitions are sampled proportionally to their learnability.
While this technique improves the sample efficiency by selecting highly-relevant samples, it does not improve the overall “learnability” of the samples themselves, and restricts training to previously observed experience.
Experience Replay Optimization (ERO) [34] parameterizes the replay buffer directly as a learned priority score function.
Rather than using heuristics such as TDerror to determine a prioritization of samples, as is performed in PER [23], in ERO this prioritization is learned directly via a policy gradient approach in which return is measured by the agent’s policy improvement [34].
A REINFORCE-based [31] estimate of the policy gradient updates the learned replay buffer in an alternating fashion with the policy being trained.
Similarly to PER, while ERO improves sample efficiency by selecting samples with high “learnability”, it does not improve the overall “learnability” of the samples themselves, and also restricts training to the agent’s observed experiences.
Interpolated Rewards Replay [28] performs linear interpolation of experienced rewards.
In contrast, NMER interpolates entire transitions using stochastic, locally linear, convex interpolation, thus allowing for richer interpolation of an agent’s experiences.
Data augmentation Data augmentation is another class of techniques used to improve the performance of DRL agents.
Reinforcement learning with Augmented Data (RAD) is a module designed for improving agent performance in visual and propioceptive DRL tasks [12].
For continuous control environments, such as the OpenAI Gym [3, 27] tasks considered in this paper, RAD leverages data augmentation techniques such as random amplitude scaling.
While this does allow for learning beyond an agent’s observed set of experiences, augmentations such as randomized amplitude scaling make no use of the other transitions stored in the replay buffer to determine amplitude scaling factors, and do not combine the existing transitions that agents store.
In Data-Regularized Q (DrQ) learning, generalizable and image-invariant data augmentation mechanisms are applied to off-policy DRL algorithms to improve sample efficiency in visual control tasks, providing off-policy agents with sample efficiency comparable to state-of-the-art MB-DRL algorithms.
Similarly to DrQ, NMER improves sample efficiency via regularization through training agents on augmented samples.
Mixup sampling Mixup was originally applied to supervised and unsupervised machine learning domains, and empirically improves the generalizability and out-of-sample predictive performance of learners [35].
Several reinforcement learning methodologies make use of Mixup-interpolated experiences for training reinforcement learning agents.
In Continuous Transition [16], temporally-adjacent transitions are interpolated with one another using Mixup, creating synthetic transitions between any two consecutive transitions.
In MixReg [29], interpolated transitions are formed using Mixup on combinations of input and output signals.
In S4RL [26], interpolated transitions are produced by interpolating current and next states within an observed transition.
While these approaches increase the training domain via interpolation, they do not strictly enforce transition locality of the resulting samples.
Proximity between the points used for sampling is encoded temporally, as in [16, 26], but not in the transition space of the agent’s experience.
NMER employs a nearest neighbor heuristic to encourage transition pairs for Mixup to be located approximately within the same dynamics regimes in the transition manifold.
Compared to Continuous Transition [16] and S4RL [26], samples interpolated with NMER may better preserve the local dynamics of the environment and enable further agent regularization through cross-episode interpolation between transitions and their dynamic sets of nearest neighbors.
Preliminaries
Neighborhood Mixup Experience Replay (NMER) builds on experience replay for off-policy DRL,
Mixup, and nearest neighbor heuristics to encourage approximately on-manifold interpolation.
Off-policy DRL for continuous control tasks Off-policy DRL has successfully been applied to continuous control tasks through the use of actor-critic methods such as Soft Actor-Critic (SAC), Deep
Deterministic Policy Gradients (DDPG), and Twin Delayed Deep Deterministic Policy Gradients (TD3) [10, 15, 7].
In this off-policy, MF-DRL setting, agents are trained using transitions composed of states, actions, rewards, and next states.
The state ( S ), action ( A ), and reward ( R ) spaces in which these transitions lie are continuous.
Experience replay.
Experience replay [18] enables an agent to train on past observations.
It can be largely decoupled from the agent’s training algorithm - while the agent seeks to learn optimal policies and value functions given observed training samples, regardless of the samples provided to it, the experience replay buffer is tasked with providing the agent samples that offer the greatest “learnability” for improving these policies and value functions.
Current experience replay approaches are discussed in Section 2.
Mixup.
Mixup [35] is a novel stochastic data augmentation technique used to improve the generalizability of deep neural networks by training on convex linear combinations of existing points.
Intuitively, this linear interpolation mechanism invokes the Bayesian linear prior on the learner that linear combinations of features result in the same linear combinations of targets [35], leading to generalizable out-of-sample performance in accordance with Occam’s Razor [20].
Through another lens, Mixup increases the generalizability of the model it is applied to by increasing the support of the model through training on interpolated examples.
Mixup performs interpolation by sampling from the convex unit line connecting two samples [35] according to a beta distribution parameterized by a hyperparameter α , which controls the spread of the distribution along this relative unit line.
To interpolate a new sample x interpolated using two existing samples x 1 , x 2 ∈ R d , Mixup interpolates according to:
[Equation]
Where α > 0 is a hyperparameter controlling the spread of the beta distribution.
On-manifold interpolation.
To measure the accuracy of interpolation in interpolated experience replay approaches, we consider how “on-manifold” the interpolated transition is with respect to the transition manifold mapping states and actions to rewards and next states.
More rigorously, we consider that the observed transitions of a replay buffer lie on a transition manifold, which we denote by the space T .
The space T is given by the Cartesian product of the state ( S ), action ( A ), reward ( R ), and next state ( S ) spaces of the agent:
[Equation]
With NMER, the use of local Mixup (convex) interpolation between only nearest neighbors acts as a heuristic to encourage interpolated samples to stay near the underlying transition manifold.
Figuredepicts on and off-transition manifold interpolation when generating interpolated transitions using
Mixup.
See our technical report for details on on-manifold assessment.
Neighborhood Mixup Experience Replay (NMER)
NMER trains off-policy DRL agents using locally linear convex combinations of an agent’s existing experiences, in effect creating locally linear models centered around each transition of the replay buffer.
By only interpolating proximal transitions with one another, where proximity is measured by the standardized Euclidean distance in the state-action space of the replay buffer, NMER interpolates transitions that have similar state and action inputs, but potentially different reward and next state outputs.
In considering these nearest neighbors, NMER regularizes the off-policy reinforcement learning agents it trains by allowing for cross-episodic interpolation between proximal transitions.
Furthermore, in the presence of stochasticity in the transition manifold, NMER can prevent RL agents from overfitting to a particular (reward, next state) outcome by interpolating different (reward, next state) outcomes for near-identical (state, action) inputs.
With NMER, the agent learns to consider not only one combination of (reward, next state) pairs given an input (state, action) pair, but rather unbiased linear combinations of them, results in the agent learning the conditional expectation of rewards and next states given states and actions.
For MF, off-policy DRL agents and continuous control tasks, the role of NMER can be visualized in Figure 2.
The steps of NMER are as follows:
Update step : When a new environment interaction is added to the replay buffer, restandardize the states and actions of the stored transitions in the replay buffer, and update the nearest neighbor data structures using Euclidean distances over the Z-score standardized, concatenated state-action features of the replay buffer.
Similarity search is thus measured over the input state and action spaces; however, we emphasize that NMER can admit other distance functions and representations of similarity as well.
See our technical report for further details.
Sampling Step : First, we sample a batch of “sample transitions” uniformly from the replay buffer.
Next, we query the nearest neighbors of each sampled transition in this sampled batch.
Following this, for each set of neighbors in the sampled training batch, we sample a neighbor transition uniformly from this set of neighbors.
Finally, for each transition in the training batch, we apply Mixup to linearly interpolate the selected samples and neighbors ( x sample ,i and x neighbor ,i , respectively) using local Mixup:
[Equation]
These steps are given in Algorithm 1.
For run time efficiency, we implement this in a batched, vectorized fashion over the training batches.
NMER introduces minimal computational overhead compared to vanilla experience replay - the only additional operations required are standardization, nearest neighbor querying, and local Mixup.
This minimal overhead positions NMER as a viable experience replay buffer for continuous control environments of any complexity and dimensionality.
Agent regularization via linear interpolation.
Through the lens of Occam’s Razor [20], NMER improves the generalizability of the policy and value function approximators of the agent by invoking the prior that linear combinations of state-action pairs result in the same linear combinations of corresponding reward-next state pairs.
This prior, although strong in some continuous control settings, improves generalizability in tasks where this linearity assumption approximately holds.
Since the spaces of an agent in continuous control tasks are continuous, interpolating continuous, linear combinations of transitions can still yield interpolated samples that lie on, or approximately on, the underlying transition manifold T .
Furthermore, if the transition manifold T is convex, NMER guarantees on-manifold interpolation, since this technique generates strictly convex combinations of transitions.
Crucially, syntheticallygenerated on-manifold transitions are indistinguishable from transitions generated at the same point using the underlying environment dynamics.
However, for many applications, particularly highdimensional, real-world continuous control tasks, the underlying transition manifold will generally be non-convex.
Neighborhood Mixup as a heuristic to encourage on-manifold interpolation.
Non-convexity and nonlinearity in continuous control environments provide motivation for our neighborhood-based interpolation mechanism, which addresses issues with non-convexity of the transition manifold by only considering interpolation between transitions in the same “neighborhood”, i.e. transitions with
Input : Replay buffer B , Mixup α > 0 , Batch Size T
Output : Interpolated Training batch B train
B = { ( s t , a t , r t , s ′ t ) }
Tt =1 iid ∼ U
( B ) //
Sample Batch Uniformly From Replay Buffer
B train ← Array [ · · · ] for t in T do ( s s , a s , r s , s ′ s ) ← B [ t ] //
Sample Transition to Interpolate With [˜ s s , ˜ a s ] T ← ZScore ([ s s , a s ] T ) //
Standardize States and Actions k s ← NN (cid:16) [˜ s s , ˜ a s ] T , B (cid:17) //
Standardized Nearest Neighbors of Sample Point ( s n , a n , r n , s ′ n ) ∼ U ( k s ) //
Sample Uniformly From Local Neighborhood x s ← [ s s , a s , r s , s ′ s ] T //
Sample Transition Features x n ← [ s n , a n , r n , s ′ n ] T // Neighbor Transition Features λ ∼ β
( α, α ) //
Sample Mixup Coefficient Using Hyperparameter α x i
= λ x s + (1 − λ )
x n // Interpolate Sample and Neighbor Transitions Using Mixup
B train [ t ] ← x i // Add Interpolated Sample to Training Batch end for return B train similar state-action pairs.
If the transition manifold is locally Euclidean, linearly interpolating two transitions is a suitable, approximately on-manifold mechanism for interpolating between spatially proximal transitions.
Continuous control evaluation
To rigorously evaluate and quantify the improvement in sample efficiency with NMER, we compare
NMER to other state-of-the-art replay buffers by applying these replay buffers to continuous control tasks and MF, off-policy DRL algorithms.
Testing environment and configuration.
We consider continuous control environments from the
OpenAI Gym MuJoCo [3, 27] suite, namely: (i) Ant , (ii) HalfCheetah , (iii) Hopper , (iv) Swimmer , and (v) Walker2d .
Since we are principally interested in evaluating the sample efficiency of replay buffers, we treat the MF, off-policy DRL algorithms used for these evaluations (SAC [10] and TD [7]) as part of the experimental configuration.
For each replay buffer variant, including NMER, we train agents using replay ratios (ratio of gradient steps to environment interactions) of 1, 5, and 20, and report the best results for each replay buffer variant.
Additionally, for SAC [10], to stabilize the policy, we add a small L2 regularizer to the actor network for all SAC evaluations.
Implementation details and ablation studies for each replay buffer variant are provided in the technical report.
We measure replay buffer sample efficiency using the evaluation reward of the reinforcement learning agent after 200K environment interactions have been sampled, as in [16, 13].
Rewards are smoothed using an averaging window of 11, as in [1].
Baselines.
We compare NMER to the following baselines: (i) Uniform, Vanilla Replay ( U ) [18, 5], where transitions are sampled i.i.d.
uniformly from the replay buffer, (ii) Prioritized Experience
Replay ( PER ) [23] with stochastic prioritization, (iii) Continuous Transition ( CT ) [16].
Since the main comparison between NMER and CT is how samples are selected for interpolation, we make two modifications to the original CT baseline: (a) We remove the automatic Mixup α parameter tuning mechanism, and (b) If a terminal state is encountered in either the sample or neighbor transition, no interpolation occurs, and the sampled transition is simply used for training the agent.
In Continuous
Transition, terminal transitions in an episode can be interpolated with their previous, non-terminal transitions, resulting in non-binary termination signals.
In order to make the fairest comparison possible between NMER and CT, we adopt the same interpolation rules we use for NMER in CT.
(iv) SUNRISE Baselines .
Additionally, we compare the performance of NMER and the baselines we implement to SUNRISE [13] and other baselines tested using the same continuous control environments and 200K environment interaction evaluation [13].
Results for these baselines are provided in Table 2.
Relative comparisons of NMER to these and additional baselines (see technical report) can be found in Table 3.
Learning curves for TD3 and SAC agents and NMER , U , PER , and CT baselines are depicted in Figures 3 and 4 ( SUNRISE and other baselines evaluated in [13] are depicted with dashed horizontal lines indicating mean evaluation reward at 200K environment interactions), and comparative tabular results after 200K environment interactions are provided in Table 1.
Each result is averaged over four runs.
Note the following abbreviations in Table 1: (i) U = Uniform Replay, (ii) PER =
Prioritized Experience
Replay, (iii) CT = Continuous Transition, (iv)
NMER =
Neighborhood Mixup Experience Replay.
Best results for each off-policy algorithm ( TD3 , SAC ) are bolded.
The results of this continuous control evaluation study indicate that NMER frequently achieves comparatively better sample efficiency than the baseline replay buffers used in this study across SAC and TD3, as well as other baseline DRL algorithms evaluated in [13]
[Table]
[Table] [Table] [Table]
Discussion
These evaluation studies demonstrate that agents trained using NMER can learn robust policies using fewer environment interactions compared to agents trained using state-of-the-art replay buffers for continuous control.
We consider the implications of NMER and its extensions to DRL.
Limitations.
Despite the observed empirical success, NMER exhibits several limitations.
NMER assumes the underlying transition dynamics are locally linear, which may be a naive approximation for continuous control tasks with nonlinear underlying dynamics.
Additionally, although NMER’s use of nearest neighbors serves as a viable heuristic to steer the replay buffer toward on-manifold interpolation, on-manifold interpolation is not analytically guaranteed.
The sections below aim to put these limitations into context by suggesting viable generalizations and considerations for future work.
Generalizing Synthetic Training.
Since NMER uses convex interpolation to generate transitions for training, as more samples are added to the replay buffer, the interpolated transitions will become more accurate.
Furthermore, NMER can be extended to modulate the ratio of real to synthetic training samples over time, enabling a variety of flexible training schemes for different replay buffer densities.
Generalizing neighborhoods.
NMER computes nearest neighbors using standardized Euclidean norms over concatenated state-action space, which allows for Mixup-based interpolation of transitions with proximal state-action vectors, regardless of the rewards and next states in these transitions.
However, this notion of proximity between transitions can be generalized to any measure of distance in the transition space of a replay buffer.
Generalizing this notion of proximity between stored environment interactions can be invoked via different distance metrics, e.g. Mahalanobis distance, as well as the use of composite product norms over different features in the transition space.
For instance, the use of a composite product norm over states × actions results in nearest neighbors having similar (state, action) pairs.
The efficacy of different proximity representations for neighborhood-based interpolated experience replay remains an open research question.
Generalizing interpolation.
NMER also invokes the implicit prior that linear combinations of states and actions result in the same linear combination of rewards and next states through the use of
Mixup.
However, under some dynamics regimes and continuous control environments, local linear interpolation via local Mixup may result in interpolated samples far from the underlying transition manifold.
For interpolating on-manifold samples in locally nonlinear neighborhoods, off-policy DRL agents may benefit from the use of more sophisticated neighborhood-based interpolation mechanisms, such as Gaussian Process Regression [19, 21] or Graph Neural Networks [36].
Conclusion
We propose Neighborhood Mixup Experience Replay (NMER), an experience replay module that improves the sample efficiency of off-policy DRL agents through synthetic sample generation.
We empirically demonstrate that training agents on experiences generated via local Mixup in the transition space of a replay buffer facilitates learning robust policies using fewer environment interactions.
NMER combines the benefits of learning from locally linear approximations of the underlying environment model with the sample efficiency benefits of learning from synthetic samples, thus expanding the possibilities for tractable DRL in real-world continuous control settings.
Acknowledgements : This research was supported by the Toyota Research Institute (TRI).
This article solely reflects the opinions and conclusions of its authors and not TRI, Toyota, or any other entity.
We thank TRI for their support.
The authors thank the MIT SuperCloud and Lincoln Laboratory
Supercomputing Center for providing HPC and consultation resources that have contributed to the research results reported within this publication