ASTRO: Adaptive Stitching via Dynamics-Guided Trajectory Rollouts

TL;DR

ASTRO introduces a novel data augmentation framework for offline reinforcement learning that generates dynamics-consistent trajectories through adaptive stitching, significantly improving policy performance on standard benchmarks.

Contribution

The paper presents ASTRO, a new method that adaptively stitches trajectories using dynamics-guided planning, addressing limitations of existing augmentation techniques in offline RL.

Findings

ASTRO outperforms prior augmentation methods on OGBench and D4RL benchmarks.

The method achieves significant policy performance improvements.

ASTRO effectively generates feasible, dynamics-consistent trajectories for offline RL.

Abstract

Offline reinforcement learning (RL) enables agents to learn optimal policies from pre-collected datasets. However, datasets containing suboptimal and fragmented trajectories present challenges for reward propagation, resulting in inaccurate value estimation and degraded policy performance. While trajectory stitching via generative models offers a promising solution, existing augmentation methods frequently produce trajectories that are either confined to the support of the behavior policy or violate the underlying dynamics, thereby limiting their effectiveness for policy improvement. We propose ASTRO, a data augmentation framework that generates distributionally novel and dynamics-consistent trajectories for offline RL. ASTRO first learns a temporal-distance representation to identify distinct and reachable stitch targets. We then employ a dynamics-guided stitch planner that adaptively generates connecting action sequences via Rollout Deviation Feedback, defined as the gap between target state sequence and the actual arrived state sequence by executing predicted actions, to improve trajectory stitching's feasibility and reachability. This approach facilitates effective augmentation through stitching and ultimately enhances policy learning. ASTRO outperforms prior offline RL augmentation methods across various algorithms, achieving notable performance gain on the challenging OGBench suite and demonstrating consistent improvements on standard offline RL benchmarks such as D4RL.