Generalized Decision Transformer for Offline Hindsight Information Matching

Overview

How to extract as much learning signal from each trajectory data has been a key problem in reinforcement learning (RL). Recent works have shown that conditioning on future trajectory information -- such as future states in hindsight experience replay (HER) or returns-to-go in Decision Transformer (DT) -- enables efficient learning of context-conditioned policies. We demonstrate that all these approaches are essentially doing hindsight information matching (HIM) -- training policies that can output the rest of trajectory that matches some statistics of future state information. We introduce Generalized Decision Transformer (GDT) framework and show how different choices for the feature function Φ(s, a) and the anti-causal aggregator not only recover DT as a special case, but also lead to novel Categorical DT (CDT) and Bi-directional DT (BDT) for matching different statistics of the future information offline.

Experiments & Video

One of the variants: Categorical DT can choose any Φ as long as the binning is tractable (e.g. reward, xyz-velocities, etc) and solve offline multi-task state-marginal matching problem. We provide the following two benchmark problems and results.

Synthesized Bi-modal Distribution

While it was trained with uni-modal trajectories (cheetah running forward, or backflipping; top two gifs), CDT successfully matches to the patchworked bi-modal distribution: cheetah running forward first and then backflipping during a single rollout (down three gifs). We term it "distribution stitching".

Diverse Unseen Distribution

When it is trained with a diverse enough dataset, CDT also can match the target distributions unseen during training. In cheetah-velocity task from meta RL/IL problem, CDT enables the cheetah to run with unknown target x-velocity (x-vel = 0.5, 1.5, 2.5).