Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Temporal Distance-aware Transition Augmentation for Offline Model-based Reinforcement Learning

Dongsu Lee, Minhae Kwon

TL;DR

TempDATA tackles offline RL in sparse-reward, long-horizon settings by learning a temporal distance-aware latent representation and a latent dynamics model to augment transitions in latent space. The autoencoder enforces macro- and micro-level temporal coherence so that latent distances reflect shortest-path costs to goals, while a latent forward model enables safe, informative rollouts used to train a goal-conditioned offline policy with an intrinsic, distance-based reward. Empirically, TempDATA outperforms prior offline MBRL methods on AntMaze variants and multi-goal Kitchen/CALVIN tasks, rivals GCRL baselines on several benchmarks, and extends effectively to pixel-based observations and dense-reward tasks. This approach reduces reliance on online data and ensembles, offering a scalable, generalizable mechanism for long-horizon planning in offline reinforcement learning."

Abstract

The goal of offline reinforcement learning (RL) is to extract a high-performance policy from the fixed datasets, minimizing performance degradation due to out-of-distribution (OOD) samples. Offline model-based RL (MBRL) is a promising approach that ameliorates OOD issues by enriching state-action transitions with augmentations synthesized via a learned dynamics model. Unfortunately, seminal offline MBRL methods often struggle in sparse-reward, long-horizon tasks. In this work, we introduce a novel MBRL framework, dubbed Temporal Distance-Aware Transition Augmentation (TempDATA), that generates augmented transitions in a temporally structured latent space rather than in raw state space. To model long-horizon behavior, TempDATA learns a latent abstraction that captures a temporal distance from both trajectory and transition levels of state space. Our experiments confirm that TempDATA outperforms previous offline MBRL methods and achieves matching or surpassing the performance of diffusion-based trajectory augmentation and goal-conditioned RL on the D4RL AntMaze, FrankaKitchen, CALVIN, and pixel-based FrankaKitchen.

Temporal Distance-aware Transition Augmentation for Offline Model-based Reinforcement Learning

TL;DR

TempDATA tackles offline RL in sparse-reward, long-horizon settings by learning a temporal distance-aware latent representation and a latent dynamics model to augment transitions in latent space. The autoencoder enforces macro- and micro-level temporal coherence so that latent distances reflect shortest-path costs to goals, while a latent forward model enables safe, informative rollouts used to train a goal-conditioned offline policy with an intrinsic, distance-based reward. Empirically, TempDATA outperforms prior offline MBRL methods on AntMaze variants and multi-goal Kitchen/CALVIN tasks, rivals GCRL baselines on several benchmarks, and extends effectively to pixel-based observations and dense-reward tasks. This approach reduces reliance on online data and ensembles, offering a scalable, generalizable mechanism for long-horizon planning in offline reinforcement learning."

Abstract

The goal of offline reinforcement learning (RL) is to extract a high-performance policy from the fixed datasets, minimizing performance degradation due to out-of-distribution (OOD) samples. Offline model-based RL (MBRL) is a promising approach that ameliorates OOD issues by enriching state-action transitions with augmentations synthesized via a learned dynamics model. Unfortunately, seminal offline MBRL methods often struggle in sparse-reward, long-horizon tasks. In this work, we introduce a novel MBRL framework, dubbed Temporal Distance-Aware Transition Augmentation (TempDATA), that generates augmented transitions in a temporally structured latent space rather than in raw state space. To model long-horizon behavior, TempDATA learns a latent abstraction that captures a temporal distance from both trajectory and transition levels of state space. Our experiments confirm that TempDATA outperforms previous offline MBRL methods and achieves matching or surpassing the performance of diffusion-based trajectory augmentation and goal-conditioned RL on the D4RL AntMaze, FrankaKitchen, CALVIN, and pixel-based FrankaKitchen.
Paper Structure (21 sections, 2 theorems, 21 equations, 9 figures, 3 tables, 1 algorithm)

This paper contains 21 sections, 2 theorems, 21 equations, 9 figures, 3 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Related Works
  3. Preliminaries
  4. TempDATA: Temporal Distance-aware Transition Augmentation
  5. Temporal Distance-aware Representation
  6. Latent Dynamic Transition Model
  7. Model-based Policy Training
  8. Algorithm Summary
  9. Experiments
  10. Performance Comparison with MBRL and TARL
  11. Performance Comparison with GCRL
  12. Performance on Pixel-based Tasks
  13. Generalizability on Dense Reward Tasks
  14. Ablation Study
  15. Scalability on Arbitrary Goal
  16. ...and 6 more sections

Key Result

Proposition 4.1

Generally, an optimal goal-conditioned value function for a state $s$ and goal $s_{\mathrm{goal}}$ is same with an optimal temporal distance as follows: where $f(s;\theta)$ is a representation encoder that captures the temporal distance of an MDP on state space $\mathcal{S}$.

Figures (9)

  • Figure 1: Performance comparison overview. (a) Umaze environment, the most naive level among AntMaze. The 8-DoF ant robot navigates the maze to reach the goal state, marked as a yellow star. (b) Comparison between the proposed solution and previous MBRL on two D4RL benchmark datasets. TempDATA (proposed) achieves the best performance in two benchmarks.
  • Figure 2: Illustration for the proposed MBRL framework. (a) Train autoencoder using offline dataset. (b) Train a dynamic model using a trained encoder and offline dataset. (c) Generate transition dataset using autoencoder and offline dataset. This happens in a representation space. (d) Extract a policy from offline and generated datasets using offline RL algorithm. Processes (c) and (d) are performed iteratively together.
  • Figure 3: Intuition of state abstraction. (Top) Temporal-aware autoencoder, our main idea, maps into state space into a representation that preserves temporal distance information. (Down) An empirical result shows that our encoder can map raw state space as latent state space in antmaze-medium environments. Here, we consider an ant position state of raw state space as fixed and use t-SNE to plot latent state space visually.
  • Figure 4: Transitions in a representation space $\mathcal{Z}$. Dashed lines are related to representation space.
  • Figure 5: Selected experimental environments. (a-c) State-based, single goal-reaching, and long-horizon navigation. (d-e) State-based, multi-goal subtasks, and long-horizon manipulation. We reuse the Kitchen environment as a Pixel-based task.
  • ...and 4 more figures

Theorems & Definitions (4)

  • Proposition 4.1: Value-metric Equivalence
  • Theorem 4.2
  • proof
  • proof