Improving Generalization and Robustness in Offline Reinforcement Learning via Boundary-Aware Data Augmentation

Gong Gao; Weidong Zhao; Xianhui Liu

cs.LGAug 22, 2026

Improving Generalization and Robustness in Offline Reinforcement Learning via Boundary-Aware Data Augmentation

Authors: Gong Gao, Weidong Zhao, Xianhui Liu

Organizations: School of Computer Science, Tongji University, China

Abstract

Current offline reinforcement learning (ORL) algorithms tend to overfit the training dataset and exhibit poor in-distribution generalization and robustness performance when deployed to real environments, thus compromising their effectiveness. Existing methods typically enhance in-distribution generalization and robustness by leveraging regularization techniques widely used in computer vision. However, due to the high sensitivity of low-level physical signals to distributional shifts, these methods still suffer from notable limitations in in-distribution generalization and robustness, making it difficult to achieve stable performance in complex environments. To address this issue, we theoretically analyze the error bounds of the behavior policy and action-value function trained with random episode interpolation, revealing that the error scales positively correlated with the distance between states. Based on this insight, we propose a method called $\bf{B}$ oundary- $\bf{A}$ ware $\bf{D}$ ata $\bf{A}$ ugmentation (BADA), which leverages neighboring states to construct interpolation boundaries, enabling the generation of synthetic data that more faithfully preserves the original data distribution. We first conduct qualitative studies in a toy environment, showing that BADA generates mixed samples that preserve desirable policy smoothness while accurately reconstructing multimodal value distributions. Extensive experiments on limited offline datasets further demonstrate that BADA attains state-of-the-art performance across diverse benchmarks.

Figures & tables

Fig. 1 : Comparative analysis of SBADA and Mixup data augmentation techniques for the walker2d-medium dataset. (a) Offline episode data $\rm{e1}$ and $\rm{e2}$ are employed as references to generate synthetic data $\mathrm{e}^{\star}$ . (b) Mixup [ 42 ] generates the first synthetic episode by interpolating two randomly selected episode pairs. A global nearest neighbor search is then applied to identify the closest states $\rm{e1}^{\textit{s}_{\textit{t}}}$ and $\rm{e2}^{\textit{s}_{\textit{t}}}$ . The second and third pairs of instances are generated using SBADA. (c) PCA [ 8 ] comparing the distribution of the synthetic data generated by vanilla Mixup and SBADA. All synthetic samples are obtained via $\mathrm{e}^{\star}=\mathrm{e1}+\alpha(\mathrm{e2}-\mathrm{e1}),\alpha=0.5$ . Consequently, synthesized data between two closely related states will not deviate from the original data distribution, whereas data generated by Mixup will result in a distributional shift.

Fig. 2 : Temporal and spatial views of an offline trajectory dataset. (a) State neighborhood retrieval based on temporal adjacency within the same trajectory, where local neighborhoods are constructed from temporally consecutive states. (b) States neighborhood retrieval according to geometric proximity in the state space, allowing neighborhoods to span across different trajectories. This illustrates that offline trajectories admit distinct neighborhood structures when viewed from temporal versus spatial perspectives.

Fig. 3 : Visualization of action distributions and state value distributions for synthetic data generated by Mixup and SBADA from sampled static data. (a) 900 offline episodes are randomly sampled to generate mixed episodes, illustrating distribution differences compared to the offline data. From left to right: statically sampled episodes, episodes generated by vanilla Mixup, and episodes generated by our SBADA. (b) Visualization of value distributions generated by Mixup and SBADA. From left to right: state values from statically sampled data, state values generated by vanilla Mixup, and state values generated by our SBADA.

Fig. 4 : Images of the Gym environments used in our experiments, including 3 continuous control tasks: (a) Halfcheetah. (b) Hopper. (c) Walker2d.

Method	10% (100K samples)									100%
Method	ANQ	Mixup	M-Mixup	RAD- $\mathcal{U}$	S4RL- $\mathcal{N}$	K-Mixup	PAN	TBADA	SBADA	ANQ
halfcheetah-r	21.7	2.2 $\pm$ 0.2	26.0 $\pm$ 0.5	19.5 $\pm$ 0.3	22.7 $\pm$ 1.1	19.4 $\pm$ 0.7	20.2 $\pm$ 0.7	22.7 $\pm$ 0.8	23.0 $\pm$ 1.2	25.1
hopper-r	14.3	3.6 $\pm$ 0.3	14.5 $\pm$ 5.7	17.5 $\pm$ 5.7	8.1 $\pm$ 2.8	4.8 $\pm$ 2.0	4.2 $\pm$ 2.0	19.6 $\pm$ 7.1	21.4 $\pm$ 5.9	31.3
walker2d-r	-0.2	1.9 $\pm$ 1.1	-0.3 $\pm$ 0.0	-0.2 $\pm$ 0.0	-0.2 $\pm$ 0.0	-0.2 $\pm$ 0.0	-0.2 $\pm$ 0.5	2.6 $\pm$ 2.8	-3.1 $\pm$ 3.0	-0.2
halfcheetah-m	40.5	41.6 $\pm$ 0.4	9.8 $\pm$ 8.5	37.5 $\pm$ 6.3	20.0 $\pm$ 8.6	-0.1 $\pm$ 0.6	30.5 $\pm$ 13.2	41.8 $\pm$ 7.9	40.7 $\pm$ 8.3	61.3
hopper-m	66.4	52.4 $\pm$ 1.2	0.7 $\pm$ 0.0	53.9 $\pm$ 3.5	89.7 $\pm$ 3.5	51.0 $\pm$ 2.7	79.9 $\pm$ 4.0	77.5 $\pm$ 7.4	85.3 $\pm$ 5.4	94.2

TABLE I : Performance comparison of TBADA, SBADA, and vanilla baselines on the Gym tasks of the D4RL benchmark. We report the mean and standard deviation of normalized D4RL scores over the final 10 evaluations and 5 random seeds, with the best scores highlighted in bold. Dataset abbreviations: r = random, m = medium, mr = medium-replay, me = medium-expert.

Method	Specification	halfcheetah-medium				hopper-medium
	Proportion	1%	2%	4%	100%	1%	2%	4%	100%
	training samples	9990	19980	39960	999000	10000	20000	40000	1000000
TD3BC	Vanilla	2.2	31.7	47.2	48.0	14.2	34.1	42.0	59.0
	+TBADA	3.9 $\pm$ 0.3	39.0 $\pm$ 0.8	48.2 $\pm$ 0.2	47.8 $\pm$ 0.1	19.7 $\pm$ 1.1	48.6 $\pm$ 3.0	50.6 $\pm$ 2.3	58.9 $\pm$ 1.9
	+SBADA	2.8 $\pm$ 0.6	36.0 $\pm$ 1.4	47.5 $\pm$ 0.5	48.2 $\pm$ 0.2	22.1 $\pm$ 1.4	53.6 $\pm$ 3.5	57.1 $\pm$ 2.5	59.6 $\pm$ 1.9
IQL	Vanilla	23.8	33.6	40.9	47.1	33.3	38.9	48.0	66.2

TABLE II : Performance on +TBADA, +SBADA, and vanilla algorithms with 10% D4RL datasets. The results for +TBADA and +SBADA correspond to the mean and standard deviation of normalized D4RL scores over the final 10 evaluations and 5 random seeds.

Fig. 5 : The performance of ANQ and IQL in the mixed quality datasets. Lines and shaded areas indicate the mean and standard deviation of the 5 random seeds, respectively. The gray dashed line indicates the learning target of the Q-network.

Method	Specification	$\epsilon_{s}=0.0$	$\epsilon_{s}=0.025$	$\epsilon_{s}=0.05$	$\epsilon_{s}=0.075$
TD3BC	Vanilla	90.8	55.5	42.2	37.5
	+RAD- $\mathcal{U}$	81.2 $\pm$ 3.4	53.9 $\pm$ 2.5	41.7 $\pm$ 2.3	36.3 $\pm$ 2.8
	+S4RL- $\mathcal{N}$	89.1 $\pm$ 1.7	57.1 $\pm$ 2.3	43.5 $\pm$ 2.0	37.7 $\pm$ 1.2
	+TBADA	90.6 $\pm$ 1.2	52.5 $\pm$ 3.8	42.9 $\pm$ 1.9	38.0 $\pm$ 1.3
	+SBADA	91.7 $\pm$ 1.5	56.8 $\pm$ 2.7	44.6 $\pm$ 1.4	39.1 $\pm$ 0.4
IQL	Vanilla	86.8	55.8	42.0	35.4

TABLE III : Performance of RAD-

\mathcal{U}

, S4RL-

\mathcal{N}

, TBADA, and SBADA combined with TD3BC, IQL, ANQ, and MOBILE under continuous noisy observations. Results report the mean and standard deviation of normalized D4RL scores over the final 10 evaluations, averaged across 5 random seeds.

TABLE IV : Comparison of the performance of IQL+TBADA, IQL+SBADA, and vanilla IQL in evaluating PGD, FGSM, I-FGSM with 15-step white-box attacks (a) and black-box attacks (b).

TABLE V : Comparison of the performance of IQL+TBADA, IQL+SBADA, and vanilla IQL in evaluating PGD and I-FGSM with 30-step and FGSM with random initialization white-box attacks (a) and black-box attacks (b).

TABLE VI : Performance comparison of CQL combined with the proposed TBADA and SBADA under the original transition dynamics

P

(a) and the noisy transition dynamics

P_{noise}

(b). All methods are evaluated on 100K offline datasets, where 20K transitions in

P_{noise}

are corrupted with random gaussian noise. Dataset abbreviations: c = complete, m = mixed, p = partial.

Fig. 6 : We compare the synthetic episode errors generated by different interpolation algorithms. (a) Comparison between the error of vanilla Mixup and that of the original offline data. (b) Comparison between the errors of our TBADA and SBADA with respect to the original offline data. (c) Comparison of the errors of SBADA using different nearest neighbor distance metrics.

Fig. 7 : PCA for dimensionality reduction of states and visualizing the nearest-neighbor states of the $L\infty$ metric in the dataset, (a) representing the distance between two time-continuous states from the temporal perspective, and (b) representing the distance between two nearest states from the spatial perspective. The nearest-neighbor state distances from the spatial perspective are closer than those from the temporal perspective.

Stage	Specification	ANQ	ANQ+TBADA	ANQ+SBADA
Initialization	GPU memory	0	0	2054M
Initialization	Consumption time	9s	9s	15s
Training	GPU memory	454M	456M	816M
Training	Consumption time	6.8h	7.1h	7.4h

TABLE VII : GPU memory consumption and training time introduced when integrating TBADA and SBADA into the ANQ algorithm.

Appendix figures & tables9 assets

Supplementary material from the paper’s appendix.

Appendix

Fig. 8 : Images of the Antmaze and Kitchen environments used in our experiments, including 4 continuous control tasks: (a) Antmaze-umaze. (b) Antmaze-medium. (c) Antmaze-large. (d) Kitchen. The three Antmaze tasks share the same state and action spaces, differing only in maze scale and data sampling method.

	Hyperparameter	Value
BADA	Augmentation choice	TBADA/SBADA
BADA	Augmentation intensity $\beta$	0.2
TD3BC	Behavior cloning scaling factor	2.5
	Optimizer	Adam [ 15 ]
	Critic learning rate	3e-4
	Actor learning rate	3e-4

Appendix

TABLE VIII : TD3BC+BADA hyperparameters.

		Hyperparameter	Value
BADA		Augmentation choice	TBADA/SBADA
BADA		Augmentation intensity $\beta$	0.2
IQL	Shared	Optimizer	Adam [ 15 ]
		Value learning rate	3e-4
		Critic learning rate	3e-4
		Actor learning rate	3e-4

Appendix

TABLE IX : IQL+BADA hyperparameters.

	Hyperparameter	Value
BADA	Augmentation choice	TBADA/SBADA
BADA	Augmentation intensity $\beta$	0.2
CQL	Optimizer	Adam [ 15 ]
	Critic learning rate	3e-4
	Actor learning rate	1e-4
	Alpha learning rate	1e-4

Appendix

TABLE X : CQL+BADA hyperparameters.

	Hyperparameter	Value
BADA	Augmentation choice	TBADA/SBADA
BADA	Augmentation intensity $\beta$	0.2
ANQ	Optimizer	Adam [ 15 ]
	Critic learning rate	$3\times 10^{-4}$
	Actor learning rate	$3\times 10^{-4}$ with cosine schedule
	Discount factor	0.99 for Gym, 0.995 for Antmaze

Appendix

TABLE XI : Hyperparameters of ANQ.

		Hyperparameter	Value
BADA		Augmentation choice	TBADA/SBADA
BADA		Augmentation intensity $\beta$	0.2
MOBILE	Shared	Optimizer	Adam [ 15 ]
		The number of critics	2
		Critic learning rate	3e-4
		Actor learning rate	1e-4

Appendix

TABLE XII : MOBILE+BADA shared hyperparameters.

Environment	Task	Penalty coefficient	Rollout length
Gym	halfcheetah-random	0.5	5
	hopper-random	0.1	5
	walker2d-random	2.0	5
	halfcheetah-medium	0.5	5
	hopper-medium	1.5	5
	walker2d-medium	1.0	5

Appendix

TABLE XIII : MOBILE+BADA specific hyperparameters for Gym datasets.

Method

TD3BC

IQL

ANQ

MOBILE

Vanilla

+TBADA

+SBADA

Vanilla

+TBADA

+SBADA

Vanilla

+TBADA

+SBADA

Vanilla

+TBADA

+SBADA

halfcheetah-r

11.3

12.5

\pm

0.3

13.0

\pm

0.2

11.2

11.4

\pm

0.5

12.1

\pm

0.3

25.1

28.8

\pm

4.0

25.7

\pm

0.2

39.0

40.3

\pm

0.1

42.1

\pm

0.2

hopper-r

8.1

8.5

\pm

0.1

8.6

\pm

0.3

7.9

7.7

\pm

0.0

8.2

\pm

0.0

31.3

31.0

\pm

0.2

30.8

\pm

0.1

30.9

34.4

\pm

0.2

31.3

\pm

0.1

walker2d-r

1.5

2.6

\pm

1.1

1.7

\pm

1.4

5.9

6.8

\pm

0.1

7.3

\pm

0.1

-0.2

\pm

0.0

-0.2

\pm

0.0

18.3

21.7

\pm

0.1

21.5

\pm

0.0

halfcheetah-m

48.1

47.8

\pm

0.1

48.2

\pm

0.2

47.1

\pm

0.1

47.2

\pm

0.1

61.3

63.0

\pm

0.3

61.5

\pm

1.2

74.5

75.4

\pm

1.3

75.0

\pm

1.6

hopper-m

58.9

\pm

1.9

59.6

\pm

1.9

66.2

68.1

\pm

2.7

66.3

\pm

1.9

94.2

95.7

\pm

2.6

100.2

\pm

0.7

106.5

106.8

\pm

0.4

107.2

\pm

0.9

Appendix

TABLE XIV : Performance comparison of TBADA, SBADA, and vanilla baselines on the Gym tasks of the D4RL benchmark. We report the mean and standard deviation of normalized D4RL scores over the final 10 evaluations and 5 random seeds, with the best scores highlighted in bold. Dataset abbreviations: r = random, m = medium, mr = medium-replay, me = medium-expert. See Table XV for complementary results on AntMaze tasks.

Method	TD3BC			IQL			ANQ
Method	Vanilla	+TBADA	+SBADA	Vanilla	+TBADA	+SBADA	Vanilla	+TBADA	+SBADA
ant-u	78.4	80.2 $\pm$ 8.7	81.7 $\pm$ 8.8	87.2	88.8 $\pm$ 6.1	87.6 $\pm$ 6.7	94.4	95.0 $\pm$ 0.8	95.4 $\pm$ 0.6
ant-ud	71.2	72.0 $\pm$ 6.3	73.8 $\pm$ 6.7	66.0	63.8 $\pm$ 3.7	64.0 $\pm$ 6.3	74.7	73.7 $\pm$ 3.1	79.9 $\pm$ 1.6
ant-mp	10.5	10.3 $\pm$ 0.4	10.3 $\pm$ 0.3	71.4	72.4 $\pm$ 1.9	76.2 $\pm$ 1.7	72.8	71.7 $\pm$ 2.5	71.5 $\pm$ 2.0
ant-md	3.0	3.3 $\pm$ 0.4	3.3 $\pm$ 0.6	69.7	71.4 $\pm$ 1.7	70.1 $\pm$ 1.6	50.3	50.8 $\pm$ 17.3	46.6 $\pm$ 16.7
ant-lp	0.2	0.2 $\pm$ 0.0	0.2 $\pm$ 0.0	39.5	44.3 $\pm$ 3.5	44.6 $\pm$ 4.3	51.9	49.6 $\pm$ 6.1	51.1 $\pm$ 3.5

Appendix

TABLE XV : Performance comparison of TBADA, SBADA, and vanilla baselines on the Antmaze tasks of the D4RL benchmark. We report the mean and standard deviation of normalized D4RL scores over the final 10 evaluations and 5 random seeds, with the best scores highlighted in bold. ant = antmaze, u=umaze, ud = umaze-diverse, mp = medium-play, md = medium-diverse, lp=large-play, ld=large-diverse.

Offline RL

Offline reinforcement learning learns decision-making policies from fixed datasets without interacting with the environment during training.

See all 298 papers

RAD: Retrieval High-quality Demonstrations to Enhance Decision-making
Jul 21, 2025Lu Guo, Yixiang Shan, Zhengbang Zhu +5Offline RL Generalization in RL
School of Artificial Intelligence, Jilin University, Changchun, China · School of Computer Science, Shanghai Jiao Tong University, Shanghai, China
Trajectory-Level Data Augmentation for Offline Reinforcement Learning
May 13, 2026Tobias Schmähling, Matthias Burkhardt, Tobias WindischData Augmentation Offline RL
University of Applied Sciences Kempten, Kempten, Germany.
Offline Policy Optimization with Posterior Sampling
May 8, 2026Hongqiang Lin, Dongxu Zhang, Yiding Sun +3OOD Generalization Posterior Sampling
Zhejiang University · Xi’an Jiaotong University · Institute of Automation, Chinese Academy of Sciences +1

Jul 21, 2025J/K move · Enter open · S save

RAD: Retrieval High-quality Demonstrations to Enhance Decision-making

Lu Guo, Yixiang Shan, Zhengbang Zhu, Qifan Liang+4

School of Artificial Intelligence, Jilin University, Changchun, China · School of Computer Science, Shanghai Jiao Tong University, Shanghai, China

Offline RL Generalization in RL

Open paper AskPDF arXiv

Offline reinforcement learning (RL) learns policies from fixed datasets, thereby avoiding costly or unsafe environment interactions. However, its reliance on finite static datasets inherently restricts the ability to generalize beyond the training distribution. Prior solutions based on synthetic data augmentation often fail to generalize to unseen scenarios in the (augmented) dataset. To address these challenges, we propose Retrieval High-quAlity Demonstrations (RAD) for decision-making, which innovatively introduces a retrieval mechanism into offline RL. Specifically, RAD retrieves high-return and reachable states from the offline dataset as target states, and leverages a generative model to generate sub-trajectories conditioned on these targets for planning. Since the targets are high-return states, once the agent reaches such a target, it can continue to obtain high returns by following the associated high-return actions, thereby improving policy generalization. Extensive experiments confirm that RAD achieves competitive or superior performance compared to baselines across diverse benchmarks, validating its effectiveness. Our code is available at https://anonymous.4open.science/r/RAD_0925_1-690E.

Abstract

Offline RL

Explore similar work

RAD: Retrieval High-quality Demonstrations to Enhance Decision-making

Trajectory-Level Data Augmentation for Offline Reinforcement Learning

Offline Policy Optimization with Posterior Sampling