Generalized Decision Transformer for Offline Hindsight Information Matching
Authors: Hiroki Furuta
Yutaka Matsuo
Shixiang Shane Gu
Venue: ICLR 2022 Spotlight (2022)
Keywords: Hindsight Information Matching
Decision Transformer
State-Marginal Matching
Hindsight Experience Replay
Reinforcement Learning

Abstract: How to extract as much learning signal from each trajectory data has been a key problem in reinforcement learning (RL), where sample inefficiency has posed serious challenges for practical applications. Recent works have shown that using expressive policy function approximators and 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 multi-task policies, where at times online RL is fully replaced by offline behavioral cloning (BC), e.g. sequence modeling. We demonstrate that all these approaches are doing hindsight information matching (HIM) -- training policies that can output the rest of trajectory that matches some statistics of future state information. We present Generalized Decision Transformer (GDT) for solving any HIM problem, and show how different choices for the feature function 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. For evaluating CDT and BDT, we define offline multi-task state-marginal matching (SMM) and imitation learning (IL) as two generic HIM problems, propose a Wasserstein distance loss as a metric for both, and empirically study them on MuJoCo continuous control benchmarks. Categorical DT, which simply replaces anti-causal summation with anti-causal binning in DT, enables arguably the first effective offline multi-task SMM algorithm that generalizes well to unseen (and even synthetic) multi-modal reward or state-feature distributions. Bi-directional DT, which uses an anti-causal second transformer as the aggregator, can learn to model any statistics of the future and outperforms DT variants in offline multi-task IL, i.e. one-shot IL. Our generalized formulations from HIM and GDT greatly expand the role of powerful sequence modeling architectures in modern RL.

Published as a conference paper at ICLR 2022
GENERALIZED DECISION TRANSFORMER FOR OFFLINE HINDSIGHT INFORMATION MATCHING
Hiroki Furuta
Yutaka Matsuo
Shixiang Shane Gu Google Research
The University of Tokyo
The University of Tokyo
Furuta@weblab. u-tooyoooo.. ac. jp
ABSTRACT
How extract much learning signal from each trajectory data has been a key problem in reinforcement learning (RL), where sample inefficiency has posed serious challenges for practical applications. Recent works have shown that using expressive policy function approximators, and 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 multi-task policies, where at times online RL fully replaced by offline behavioral cloning (BC), e.g. sequence modeling. We demonstrate that all these approaches are doing hindsight information matching (HIM) - training policies that can output the rest of trajectory that matches some statistics of future state information. We present Generalized Decision Transformer (GDT) for solving any HIM problem, and show how different choices for the feature function 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. For evaluating CDT and BDT, we define offline multi-task state-marginal matching (SMM) and imitation learning (IL) as two generic HIM problems, propose a Wasserstein distance loss as metric for both, and empirically study them on MuJoCo continuous control benchmarks. Categorical DTT which simply replaces anti-causal summation with anti-causal bimning in DT, enables arguably the first effective offline multi-task SMM algorithm that generalizes well to unseen (and even synthetic) multi-modal reward or ftaturee distributions. Bi-directional DT, which uses an anti-causal second transformer as the aggregator, can learn to model any statistics of the future and outperforms DT variants in offline multi-task IL. i.e. one-shot IL. Our generalized formulations from HIM and GDT greatly expand the role of powerful sequence modeling architectures in modern RL.
INTRODUCTION
Reinforcement learning (RL.) suffers from the problem of sample inefficiency, and central question is how to extract as much learning signals, or constraint equations (Pong et al., 2018; Tu & Recht, 2019; Dean et al., 2020), from each trajectory data as possible. As dynamies transitions and Bellman equation provide rich source of supervisory objectives and constraints, many algorithms combined model-free with model-based, and bolicy-based with value-based in order achieve maximal sample efficiency, while approximately preserving stable, unbiased policy learning (Heess et al., 2015; Gu al., 2016; 2017; Buckman et al.. 2018; Pong et al., 2018; Tu & Recht, 2019).
Orthogonal to these, in the recent years we have seen a number of algorithms that are derived from different motivations and frameworks, but share the following common trait: they use future trajectory information 7:r to accelerate optimization of contextual policy x <a [st, z) with context z with respect to a parameterized reward function r(st, as. z) (see Section for notations). These hindsight algorithms have enabled Q-learning with sparse rewards (Andrychowicz et al., 2017), cemporall--extent object manipulation in cluttered scenes from human play (Lynch al., 2019). efficient multi-task RL (Eysenbach et al., 2020; Li et al., 2020), offline self-supervised discovery of manipulation primitives from pixels (Chebotar al., 2021), 1), and offline RI. using return-conditooeed supervised
model-based RL with Q-function (Pong et al., 2018), mastery of 6-DoF

---

Published as a conference paper at ICLR 2022
learning with transformers (Chen al., 2021a; la; Janner et al., 2021). We derive generic problem formulation covering all these variants, and observe that this hindsight information matching (HIM) framework, with behavioral cloning (BC) the learning objective, can learn a conditional policy to generate trajectories that each satisfy any properties, including distributional.
Given this insight and recent casting of RL sequence modeling (Chen et al., 2021a; Janner et al., 2021), we propose Generalized Decision Transformer (GDT), a family of algorithms for future information matching using hindsight behavioral cloning with transformers, and greatly expand the applicability of transformers and other powerful sequential modeling architectures within RL with only small architectural changes to DT. In summary. our key contributions are:
We introduce hindsight information matching (HIM) (Section 4, Table 1) as a unifying view of existing mundsight-inspieee algorithms, and Generalized Decision Transformers (GDT) as generalization of DT for RI. as sequence modeling to solve any HIM problem (Figure 1).
Inspired by distribution RL (Bellemare al., 2017: 7; Dabney al., 2018) and state-marginal matching (SMM) (Lee et al., 2020; Ghasemipour et al., 2020; Gu et al., 2021), we define offline multi-task SMM problems, propose Categorical DT (CDT) (Section 5), validate its empirical performance to match feature distributions (even generalizing to synthetic bi-modal target distribution at times), and construct the first benchmark tasks for offline multi-task SMM.
Inspired by one-shot imitation learning (Duan et al., 2017; Finn et al., 2017; Dasari & Gupta, 2020), we define offline multi-task imitation learning (II.), propose Wasserstein-distane evaluation metric, develop Bi-directional DT (BDT) as fully expressive variant of GDT (Section 5), and demonstrate BDT's competitive performance at offline multi-task IL.
causal transformer
Method
i a)
Aggregator
DT (Chen al., 2021a)
r(s,a) Learned
Summation
anti -causal aggregator
DT-X (Section 5.3)
Summation
CDT (Section 5.2)
r(s,0) a) any
Binning
BDT (Section 5.4)
Learned
Transformer
1-1
Figure 1: Generalized Decision Transformer (GDT), where the figure a minor generalization of the DT architecture (Chen et al., 2021a) and the table summarizes how it leads to different classes of algorithms with only small architectural changes. If the feature function 4ss, a) is reward r(s, a) and the anti-causal aggregator is q-discounted summation, we recover DT for offline RL. If the aggregator is binning, we get Categorical DT (CDT) for offline multi-task state-marginal matching. the aggregator is a second transformer, we get Bi- di-rrrcctional DT (BDT) for offline multi-task imitation learning (IL), or equivalently one-shot IL. The choices of (s,, a) and the aggregator together decide (T) in Hindsight Information Matching (HIM) objective discussed in Section 4 and Table 1, where conversely GDT can essentially solve any HIM problem with proper choices of & and aggregator.
2 RELATED WORK
Hindsight Reinforcement Learning and Behavior Cloning Hindsight techniques (Kaelbling, 1993; Andrychowicz et al., 2017; Pong et al., 2018) have revolutionized off-policy optimization with respect to parameterized reward functions. Two key insights were (1) for off-policy algorithms such as Q-learning (Mnih et al., 2015; Gu et al., 2016) and actor-critic methods (Lillicrap et al., 2016; Haarnoja et al., 2018; Fujimoto et al., 2018; Furuta et al., 202 the same transition samples can be used to learn with respect to any reward parameters, as long as the reward function is re-computable, ie. "relabel"-able smooth with respect to the reward parameter, generalization can speed up learning even with respect to "unexplored" rewards. In goal- based RL where future states can inform "optimal" reward parameters with respect the transitions' actions, hindsight methods were applied successfully to
like goal reaching rewards, and (2) if policy or Q-functions
2

---

Published as a conference paper at ICLR 2022
enable effective training of goal-based Q-function for sparse rewards (Andrychowicz et al., 2017),
derive exact connections between Q-learning and classic moodel-based RL (Pong et al., 2018), data- efficient off off-policy hierarchical RL (Nachum et al., 2018), multi-task RL (Eysenbach et al., 2020; Li et al., 2020), offline RL (Chebotar et al., 2021), and more (Eysenbach et al., 2021; Choi et al., 2021; Ren al., 2019; Zhao & Tresp, 2018; Ghosh et al., 2021; Nasiriany al., 2021). Additionally, Lynch et al. (2019) and Gupta et al. (2018) have shown that often BC is sufficient for learning generalizable parameterized policies, due to rich positive examples from future states, and most recently Chen et al. (2021a) and Janner et al. (2021), when combined with powerful transformer architectures ( (Vaswani et al., 2017), produced state-of-the-ar of-the-art offline RL and goal- baaad RL results. Lastly, while motivated from alternative mathematical principles and not for parameterized objectives, future state information was also explored as ways of reducing variance or improving estimations for generic policy gradient methods (Pinto et al., 2017; Guo et al., 2021; Venuto et al., 2021).
Distributional Reinforcement Learning and State-Marginal Matching Modeling the full distri- bution of returns instead of the averages led to the development of distributional RL algorithms (Belle- mare et al., 2017; Dabney et al., 2018; 2020; Castro et al., 2018; Barth-Maron et al., 2018) such as Categorical Q-learning (Bellemare et al., 2017). While our work shares techniques such discretiza- tion and binning, these works focus on optimizing a non-conditional
wwwrd-maximiziigg RL policy and therefore our problem definition is closer to that of state-marginal matching algorithms (Hazan et al., 2019; Lee et al., 2020; Ghasemipour et al., 2020; Gu et al., 2021), ), or equivalently inverse RL algorithms (Ziebart et al., 2008; Ho & Ermon, 2016; Finn et al., 2016; Fu al., 2018; Ghasemipour et al., 2020) whose connections to fcature-expectation matching have been long discussed (Abbeel & Ng, 2004). However, those are often exclusively online algorithms even sample-efficient vari- ants (Kostrikov et al., 2019), since density-ratio estimations with either discriminative (Ghasemipour et al., 2020) or generative (Lee et al., 2020) approach requires on on-policy samples, with a rare excep- tion of Kostrikov et al. (2020). Building on the success of DT and brute-force hindsight imitation learning, our Categorical DT is to the best our knowledge the first method that benchmarks offline state-marginal matching problem in the multi-task settings.
RL and Imitation Learning as Sequence Modeling When scaled to the extreme levels of data and computing, sequence models such transformers (Vaswani et al., 2017) can train models to master an impressive range of capabilities in natural language processing and computer vision (Devlin et al., 2019; Radford et al., 2019; Brown et al., 2020; Radford et al., 2021; Ramesh et al., 2021; Chen et al., 2021b; Bommasani et al., 2021; Dosovitskiy et al., 2020). Comparing to their popularity in other areas, the adoption of transformers or architectural innovations in RL have been slow, partially due the difficulty of using transformers over temporal scales for online RI. (Parisotto et al., 2020). Recent successes have focused on processing variable-length per-timestep information such as morphology (Kurin et al., 2021), sensory information (Tang & Ha, 2021), I), one-shot or few-shot imitation learning (Dasari & Gupta, 2020), or leveraged offline learning (Chen et al., 2021a; Janner et al., 2021). Our formulation enables sequence modeling to solve novel RL problems such as state- marginal matching with minimal architectural modifications to DT, greatly expanding the impacts of transformers and other powerful sequence models in RL.
3 PRELIMINARIES
We consider a Markov Decision Process (MDP) defined by the tuple of action space A, state space S, transition probability function p(s'|s, (s, a), initial state distribution p(so), reward function r(s, a). and discount factor 7 E (0, 1]. In deep RL, policy that maps the state space the action space is parameterized by the function approximators, T The RL objective is given by:
1
LRI(TT) 1- E Be-pr (0) (o) [r(s,a)]
(1)
where (s) = sou Ront II p(s s at ( (arlss)) and (s) = (1-7) Ee T pf T s)
are short- chnds- for time-alignn aligned and time-aggregated state marginal distributions following policy π.
For simplicity of notations, we write Markovian policies; however, such notations can easily apply to non-Markov policies such as Decision Transformer (Chen et al., 2021a) by converting to augmented MDP consisting of past N states, where N is the context window of DTT.
3

---

Published as a conference paper at ICLR 2022
3.1 STATE MARGINAL MATCHING
State marginal matching (SMM) (Lee et al., 2020; Hazan et al., 2019; Ghasemipour et al., 2020) has been recently studied as alternative problem specification in RL, where instead of stationary-reward maximization, the objective is to find policy minimizing the divergence D between its state marginal distribution p² p=(ss given target distribution (s)
(TMM( (D -D(pp(s)),') (s).
where D is divergence measure such as Kullback-Leibler (KI.) divergence (Lee et al., 2020; Fu et al., 2018) or, more generally, some f-divergences (Ghasemipour et al., 2020). For the target distribution p" (s), Lee et al. (2020) set uniform distribution to enhance the exploration over the entire state space; Ghasemipour et al. (2020) and Gu et al. (2021) set through scripted distribution sketches generate desired behaviors; and adversarial inverse RL methods (Ho & Ermon, 2016 6; Fu et al., 2018; Ghasemipour et al., 2020; Kostrikov et al., 2020) set as the expert data for imitation learning. Notably, unlike the RL objective in Eq. SMM objectives like Eq.2 longer depend on task rewards and are only functions of state transition dynamics and target state distribution.
3.2 PARAMETERIZED RL OBJECTIVES
Lastly, we discuss the basis for methods like HER and TDM (Andrychowicz et al., 2017; Pong et al., 2018), LiP (Lynch et al., 2019), and return-conditionss or upside-down RL (Srivastava et al., 2019; Kumar et al., 2019; Chen et al., 202 Janner et al.. 2021): parameterized RL objectives. Given parameterized reward functions with parameter E Z, conditional policy < 2) is learned with respect to multiple values of z simultancously weighted by p(z). As examples, the RL objective in E becomes:
(T) = Ez [LRI. (T, 2)] = Ex- px (c)) (r=(s, a)
(3)
where the state marginal p is from rolling out a conditioned policy π ((11:))). These can be considered as special case of contextual MDPs (Jiang et al., 2017) and are all multi-task RL problems.
4 HINDSIGHT INFORMATION MATCHING
We show how HER and TDM (Andrychowicz et al., 2017; Pong al., 2018), LiP (Lynch et al., 2019), hindsight multi-task RL (Li et al., 2020; Eysenbach et al., 2020), and return-conding condittond or upside-down RI. (Srivastava et al., 2019; Kumar al., 2019; Chen et al., 2021a; Janner al., 2021) all belong to hindsight algorithms with a shared idea of using future state information to automatically mine for positive, or "'pptimal', examples with respect to certain contextual parameter values, where these examples can accelerate RL or be used for behavior cloning (BC), i.c. supervised learning. We start by defining additional notations.
Given a partial trajectory from state 8t as T {st, a, St-1 J.., we define its information statisties as [(F.). I(F)) could be any function of trajectory that captures some statistical properties in state-space or trajectory-space, such as sufficient statistics of a distribution. like mean, variance or higher-order moments (Wainwright
Jordan, 2008). For convenience, we further define the notion of feature function +(:.)) S x A + FS, where the trajectory is then noted as t (dr, dr+1. dT) ) di- = 4(s,, at) E and the information statistics as 1 (Tr). & in practice
AAs discussed in Fu et al. (2018) and Ghasemipour et al. (2020). it's also ttright-fooward to define state- aattonnmarginal matching with respect to p² (s, a) p² (a))=(aa)s and the exact same algorithms apply. 'There a rich literature on one-shot, or few-shot, imitation learning through meta learning. For closer connections to parameterized policies and relaheling, we mainly discuss metric-hased (or amortization- methods (Duan et al., 2016), opposed to gradient-
approaches (Finn et al., 2017).
*DisCo RL (Nasiriany et al., 2021) conditions on a parameterized goal distribution and uses hindsight techniques; however, their RL. objective in Equation (e) [log (8)), in contrast to a proper divergence objective like Ghasemipour al. (2020), missing the H (s)) entropy term and is essentially just solving for parameterized stationary reward maximization, as also stated in their Remark 1.
'Recent mutual information maximization or empowerment methods (Fysenbach et al., 2019; Sharma al., 2020; Choi et al., 2021) also make similar assumptions; Gu al. (2021) for more details.
4

---

Published as a conference paper at ICLR 2022
Method
Alga Type
Training
(r)
Architectures
Andrychorcc (2017)
RL
Online
MLP
Pong et (201))
RI.
Online
Chebotar et al. (2021)
RL.
Offline
CNN
et
Online
ang
MIP
Bysesboch ot al. (2020)
BCRL
OwNietine
arg max Lx (gy))
MLP
Tanch (2019)
BC
Offinee
Stochaotic RNN
Ghosh (2021)
Online
or
Srivastava et (2019)
BC
Online
Fist Weighes
Komar (2011)
BC
Online
Tonner et al. (2021)
BC
Offine
Transfommr
et (2017)
Offline
MLP +
Geseralized DDT (ours)
BC
Offiine
Any
Transfocmer
DI (Chen al., 2021a)
BC
Offline
Σ, >'re
Transformer
Categorical DT (ours)"
BC
Offine
histugramtre. (7)
Transfommr
Di-eiieclnnal DTT (ours)
BC
Offine
Transformer
Table coarse summary hindsight information matching (HIM) algorithms. The notation follows Scction With HIM, all prior works can categorized to four generic problem types based on (T): (1) goal- or (Andrychowicz ct al., 2017), 7) (2) multi-task arg max Σ. 7 (s, de, -) (Li ct al., 2020), (3) return-based E: 7'r, (Chen et al., 202 or (4) full trajectory imitation т (Duan et al., 2017). 4 the reward function r(s, a) in (2) and (3), an indexing function for state dimensions (e.g. xy-velocities) or a learned function (Nair ct al., 2018) in (1), or an identify function in (4). Our CDT introduces new category, (5) distribution- (r) histogram(rr, 7), based on minimal modification DT, while our BDT can be considered DT adapted for (4), the trajectory imitation.
can be an identity function, the reward function r(s. a). sub-dimensions of 8 (e.g. xy-velocities), or a generic parameterized function (e.g. auto-encoder). Generalizing reward-centric intuitions in DT (Chen et al., 2021a), we define information matching (IM) problems learning conditional policy π {als,, 2) whose trajectory rollouts satisfy some desired information statistics value 2:
min E Ex-p(()) [D(I (T].
(4)
An important observation for the IM objective (Eq.4) is that for any given trajectory т, setting z" = (7) will minimize the inner term divergence D and therefore states and actions are optimal with respect to = and samples of (n, z; can be used to accelerate RL or do BC. We call these algorithms hindsight information matching (HIM) algorithms.
Table 1, which classifies prior methods into effectively four categories based on (r), leads us to have the following insights around HIM algorithms:
New HIM algorithms can be proposed by simply changing Ist (T), we did to propose Categorical DT for (5) distribution-based.
Given a choice of _ (T). new HIM algorithms can be proposed by changing implementa- tion details (Furuta et al., 2021a), such as using "RL" or "BC" as algorithm type. doing "Online" or "Offline" training (Levine et al., 2020), and network architectures. All "Of- fline" "BC" methods could be adopted easily to "Online" learning through recursive data collections (Ghosh et al., 2021; Kumar et al., 2019; Matsushima et al., 2021).
Only (1) goal-based and (2) multi-task our (5) distribution-based can use "BC", because "RL." requires optimizing Eq. with respect to the policy, which gets non-trivial for some choices of I$t (7); e.g. (3-5) return-based, full trajectory imitation, or distribution-based. "BC" bypasses the need to solve Eq. 4 and therefore universally applicable to any I (7) or HIM algorithm
use "RL." as algorithm type, while all four, plus
GENERALIZED DECISION TRANSFORMER
Following the insights in Section 4, we introduce Generalized Decision Transformer (GDT), which generalizes DT (Chen et al., 2021a) based on different choices of (T), as described in Figure I and the last rows of Table 1. We chose DT as the base model since is simple model that uses "BC" as the algorithm type and "Transformer" as the architecture. The choice of "BC" is must, so we can tractably train GDT with respect to any 1 (T) or HIM problem. The choice for the architecture is
'Evoluioonary strategies (Salimans et al., 2017), technically nIn-RI. black-box algorithm, could be applied, but accurate estimations of D in E could require prohibitively many samples.
5

---

Published as a conference paper at ICLR 2022
more flexible; however, decided to use transformers (Vaswani et al., 2017) in this work due to their enormous scaling successes in language and vision domains (Dosovitskiy et al., 2020; Brown et al., 2020; Ramesh et al., 2021). See Algorithm 1 (in Appendix F) for the full pseudocode. While GDT in Figure 1 can lead to different algorithms depending on different choices of the feature function ((s, a) and the anti-causal aggregator (which together determine (r)). in this work we focus our empirical studies on the following two variants: Categorical DT (CDT) and Bi-directional DT (BDT).
5.1 TASK DEFINITIONS AND METRICS
Before proceeding to define CDT and BDT, we first concretely define the tasks they are designed to solve, namely: offline multi-task state-marginal matching (SMM), and offline multi-tass imitation learning (IL). Given the intrinsic connection equivalence between distribution matching and IL (Ghasemipour et al., 2020), these two separate terminologies may seem redundant. However, inspired by the initial papers studying SMM problems (Lee al., 2020; Ghasemipour et al., 2020) which qualitatively evaluates distribution matching results in specified state dimensions (e.g. xy- positions), we define the imitation task offiin multi-task SMM if specific & is given, and as offline multi-task IL if & is an identity (i.c. Ø = s) or learned (e.g. auto-encoder). We essentially view IL as SMM evaluation on full state.
Given these definition, we also define a single metric for both offline multi-task SMM/IL: typical IL assumes some availability of task reward or success evaluation, and indirectly measure the quality of imitation through it (Ho & Ermon, 2016; Fu et al., 2018). Instead, again grounding on its connection to distribution matching (Ghasemipour et al., 2020), we propose a Wasserstein loss between state- marginal and target distributions as SMM-inspired metrics for evaluating offline multi-task SMM or IL tasks. However, it is often intractable to measure such loss for full state or even for some state dimensions analytically because both state-marginal and target distributions can be non-parametric and we cannot access their densities. In practice, we empirically estimate it employing the binning of the feature space we specified. More discussions are included in Appendix C.
5.2 CATEGORICAL DECISION TRANSFORMER FOR DISTRIBUTION MATCHING
Inspired by the recent successes in distributional RL (Bellemare et al., 2017; Dabney et al., 2018; 2020), offline RL (Fujimoto et al., 2019; Jaques et al., 2020; Ghasemipour et al., 2021; Fujimoto & Gu, 2021) and state-marginal matching (Lee et al., 2020; Ghasemipour et al., 2020; Gu et al., 2021), we introduce Categorical DT (CDT) for offline state-marginal matching (SMM) problem in Section 5.1. Following the prior works (Bellemare et al., 2017; Furuta et al., 2021b), we assume oow-dimensional l e.g. rewards state dimensions like xyz-velocities, and employ the discretization of feature spaces to form categorical approximations of continuous distributions. To compute 27 = I (Ter) for all timesteps t given trajectory T1:2, we use similar recursive Bellman-like computation inspired by Bellemare et al. (2017) (see Appendix F for the details). To the best of our knowledge, this is the first paper study offline multi-task SMM and propose effective algorithm for it.
5.3 DECISION TRANSFORMER WITH LEARNED +
While CDT assumes some low-dimensional + provided for tractable binning and distribution approximation, we also study cases where & or is not provided, and instead &b is learned through auto-encoding (Hinton & Salakhutdinov, 2006; Bengio et al., 2012) (DT-AE) or contrastive (van den Oord et al., 2018; Srinivas et al., 2020; Yang & Nachum, 2021) (DT-CPC) losses for DT (see Appendix G for the details). In this case, CDT is unnecessary because if + learns sufficient features of s, matching their means, i.e. moments, through DT is enough to match any distribution to an arbitrary precision ( (Wainwright & Jordan, 2008; Li et al., 2015). Since t is differentiable with respect to DT's action-prediction losses, we also compare DT-E2E, where we learn & through end-to-end differentiation. As Section 5.1 defines, these methods do offline multi-task SMM with full state, or offline multi-task IL, similar objective to state- marginal matching or adversarial inverse RL (Ho Ermon, 2016; Ghasemipour et al., 2020) in online RL. To the best of our knowledge, this is the first offline multi-task II. method that explicitly accounts for SMM through architectural bottlenecks.
5.4 BI-DIRECTIONAL DECISION TRANSFORMER FOR ONE-SHOT IMITATION LEARNING

---

Published as a conference paper at ICLR 2022
The absence of given & could be tackled with learning not only parameterized T as in Section 5.3, but also a parameterized aggregator. Building on the connection to one-shot imitation learning (Duan et al., 2017), provide another natural extension of DT under GDT framework called Bi-directional Decision Transformer (BDT), which assumes an identity Φ and learns representation within the aggregator, in this case second (anti-causal) transformer (Radford al., 2018) that takes reverse- order state sequence as an input. See Algorithm I in Appendix F for the pscudocode, and Appendix D for further comments on the connection to one-shot or meta learning. While we found some positive results for even simple unsupervised regularizer approaches (e.g. DT-AE), we observe BDT could achieve substantially better offline multi-task IL results than DT-X variants Section 5.3.
6 EXPERIMENTS
We empirically investigate the following questions:
(SMM) Can CDT match unseen reward distributions?
(SMM) Can CDT match and generalize to unseen 1D/2D state-feature distributions?
(SMM) Can CDT match unseen synthesized state-feature distributions?
(IL) Can BDT perform offline one-shot imitation learning in full state?
We experiment on the OpenAl Gym, MuJoCo tasks (HalfCheetah, Hopper, Walker2d, Ant-v3), common benchmark for continuous control (Brockman et al., 2016; Todorov et al., 2012). Through the experiments, we use medium expert datasets D4RL (Fu et al., 2020) to ensure the decent data coverage. We sort all the trajectories by their cumulative rewards, hold out five best trajectories and five 50 percentile trajectories as a test set (10 trajectories in total), and use the rest as a train set. We report the results averaged over 20 rollouts every 4 random seed. We share our implementation to ensure the reproducibilityy..
As discussed in Section 5.1 and Appendix C, we evaluate CDT/BDTT with approximate distribution matching objective: Wasserstein-1 I distance between categorical distributions of features in rollouts or target trajectories. We compare CDT to DT, Meta-BC, and FOCAL (Li et al., 2021), metric-based offline meta RL method, as baselines. While Meta-BC and FOCAL does not solve the offline distribution matching problem directly, they provide decent baseline performance since their offline one-shot adaptation to the given target trajectories could deal with it (see Appendix B for the details).
6.1 REWARD AND STATE-FEATURE MATCHING
First, we evaluate whether CDT could match its rollout to the target distribution. We choose reward and state-feature, such as x-velocity of the agents, feature spaces to match. To specify the target distributions during the evaluation, we feed the categorical representation of the target to CDT. As the same as the reward case, DT takes the summation of the state-feature over trajectory as an input.
We quantitatively compare CDT against baselines (Table 2) in x-velocity case, where CDT shows better matching results to the target distributions unseen during training. We provide the reward distribution results and the visualization in Appendix E. where CDT performs very well in all cases. To test the generalization additionally, we intervene the target distributions by (1) shifting the target distributions in the test set by constant offsets, and (2) synthesizing novel distributions via python scripts. See Appendix E.7 and E.8 for the results. Furthermore. to investigate the scalability state-features, we experiment 2D state-feature distribution matching
of CDT to multi-dimension (xy-velocities on Ant) in Appendix E.3, where CDT outperforms other baselines.
Medied
Experi
Modian
Total
Expert
Modium
Tora
Expari
Mediun
Votal
Culepatioal
0.329
0:814
1.00
0/04
0,159
0,069
D1136
1000
UTT
0:346
10.50
1.0e
1.549
0991
1
1.177
6.81
8.07
(L135
01851
0.044
0.115
nom
0.587
BC (nsorntoo)
3.007
0.890
5.608
1.211
3:342
1 121
B882
0.264
0.548
190
84600 1 0.300
0.248
0.67
0301
0.77
1,008
Men-ai::
0.888
0340
0848
0991
0.190
04464 1.0.4
0.100
1485
1.52
0.690
Table 2: Quantitative evaluation of ftaturfe distribution matching, measuring Wasserstein- distance between the rollout and target distributions. We compare Categorical DT against DT, BC, Mcta-BC, and FOCAL. CDT achieves better matching than basclines. See Table in Appendix E. for the reward distribution results.
http://///ggtabb.cco/ttttt////// wieeeeerralidddddee ized_dt
7

---

Published as a conference paper at ICLR 2022
6.2 GENERALIZATION TO UNSEEN TARGET DISTRIBUTION
The performance of offline methods in RL is often restricted by the coverage or quality of datasets. While we demonstrate CDT can perform offiine state-magginal matching unseen target distributions in Section 1,1, the standard offline datasets might not be diverse enough to observe the generalization since those are collected by single-task to more diverse behaviors, we investigate the following tasks: (I) z-velocity distribution matching with synthesized bi-modal behavior and (2) cheetah- matching problem from meta RL/II. literature (Rakelly et al., 2019; Li et al., 2021; Fakoor et al., 2020)
wwward-maximizatee
policies. To test the generalization
6.2.1 SYNTHESIZING UNSEEN BI-MO DISTRIBUTION
To generate diverse behaviors for z-axis, we obtain the expert cheetah that backflips towards -x direction by modifying reward function (see Appendix E.4 for the details). Combining expert backflipping trajectories and expert running forward trajectories from D4RL dataset, we construct a novel dataset with diverse behaviors. We experiment the offline SMM with not only each uni-modal behavior (backflipping or running forward). but also synthesized bi-modal behavior, running forward first, then backflipping during single rollout, using patchworked target trajectories.
Table 3 and Figure 2 (a) show that CDT successfully matches the distribution to both uni-modal (running forward or backflipping) and synthesized bi-modal distributions better than DT and FOCAL, and is comparable to Meta-BC that originally designed to deal with such multi-task settings. Due to the difficulty for RL algorithms to maximize Eq. 4, FOCAL struggles to solve the offline multi-task SMM, even though FOCAL uses the same context embedding as Meta-BC. The bi-modal behavior learned by CDT can be seen at https: sites. google. com/ hwt///iiwwee/eeeeeeeeveeeeeeeee/ hee///gneraiiiee aimeeraliiddd.... zeddt.
Method
Uni-modal
Bi-modal
Average
Categorical DT
I,562 0.632
1.625 = 0,902
1.594
DT
2.676 = 0.765
2.703 + 0.703
2.690
Meta-BC
1.519 0.696
1.635 1 0.990
1.587
FOCAL (Li et al., 2021)
2.203 1 1.050
1.983 1 0.948
2.093
Table 3: State-feature (z-velocity) distribution matching with uni-modal and (synthesized) bi-modal target trajectories in HalfCheetah environment. Categorical DT matches both uni- and bi-modal trajectories better than DT and FOCAL, and is comparable to Meta-BC that originally aims to solve multi-task problem.
Eireatiivg
(a) Z-Volocity
(b) Unseen Cheetah- -elocity
Figure 2: (a) Z-Velocity and (b) Unseen Cheetah- Cheetah-Velocity results. Blue histograms represent target distributions. In (a), CDT (red) match only uni-modal behaviors for both running forward and backflipping but also bi-modal behaviors; during a single rollout running forward first, then backflipping. DT (yellow) tends to lean backtlipping and fails to fit neither uni modal nor bi- bi-modal ones. In (b), CDT successfully handles the trajectories unseen during training, while DT seems to output covering behaviors over the dataset support.
6.2.2 DIVERSE UNSEEN DISTRIBUTION FROM META LEARNING TASK
Generalization to unknown target demonstrations or tasks has been actively investigated meta or one-shot RL/IL literature (Duan et al., 2016; Wang et al., 2016). To verify the generalization of CDT to diverse behaviors, we adopt the cheetah-velocity task; popular task meta RL/IL (Rakelly et al., 2019; Li al., 2021; Fakoor al., 2020), where the cheetah tries to run with the specified velocity. We prepare 31 target x- x-velocities; taken from [0.0, 3.0], uniformly spaced at 0.1 intervals, and hold out (0.5, 1.5, 2.5) as test set. See Appendix E.5 for the dataset generation.
8

---

Published as a conference paper at ICLR 2022
Table 4 and Figure (b) reveal that CDT outperforms DT or FOCAL, and is slightly better than Meta- BC through the distribution matching evaluation, which implies CDT could solve offline multi-task SMM generalizing to the unknown target trajectories, given sufficiently diverse offline datasets.
Method
x-vel: 0.5
x-vel: 1.5
X-vel: 2.5
Average
Categorical DT
0.060 1 0.026
0,211 0.022
0.149 0.110
0.140
DT
1.197 + 0.227
0.533 1 0.105
0.861 = 0.247
0.864
Mcta-BC
0.150 = 0.069
0.152 + 0.127
0.167 = 0.055
0,156
FOCAL (Li et at. 2021)
0.472 = 0.005
0.952 = 0.073
0.346 0.186
0.590
Table 4: Generalization to unseen target velocity (x-velocity = 0.5, 1.5, among training dataset (x-velocity E [0.0, 3.0 (0.5, 1.5, 2.5]; uniformly spaced 0.1 intervals), which is a popular setting in meta RL//IL CDT successfully deals with unseen target velocities well, outperforming DT and FOCAL, and is slightly better than Meta-BC through the offline multi-task SMM evaluation.
6.3 ONE-SHOT DISTRIBUTION MATCHING FULL STATE
Lastly, we investigate BDT for offline multi-task IL. where we do not observe rewards nor state features explicitly. Instead, target full state trajectories that the agents are expected to mimic is given. We compare BDT against parameterized & variants; DT-AE, -CPC, and -E2F discussed in Section 5.3. We consider three strategies to train the encoder for DT-AE and -CPC: training with only unsupervised loss, training with unsupervised and DT's supervised loss jointly (called as "joint"), and pre-training with only unsupervised loss and freezing the weights during DT training (called as "frozen").. We aggregate the learned & by summation within the input sequence (length is 20). We evaluate BDT and DT-X with the same distribution matching evaluation as Section 6.1.
We sweep m-dim learned feature from the encoder (s)) with m E {1,4, 4, 16). Table 5 presents the offline multi-task IL results with respect to x-velocity distribution, using m = 16 embeddings (see Appendix E.6 for the full results including m 1, 4, and the reward distribution case). Similar the discussion in Yang & Nachum (2021), DT-CPC sometimes fails to obtain sufficient representation for imitation. While even simple approaches, DT-AE and DT-AE (frozen), show positive results compared to no-context BC baselines (in Table 2), BDT outperforms all other learned & variants or Meta-BC and is comparable CDT or DT (also in Table 2) with longer input (NV 50). This implies that even though we don't 't assume the feature specification, aggregator choice in GDT with minimal architectural changes may solve the offline distribution matching problem efficiently. We leave more sophisticated objectives or architectural changes as future work.
Method
Expert
hallchectah
Mediun
Total
Expert
Malisci
Toxa
Expent
wolkerS
Mediun
Total
DT-AE
20050 2: 1106
1800 2: 2U22
1,534
1.05
DSTI 2 C0000
0.142 E: (.B4)
0.377
0:341
0,833 = 0.346
[L311 0.122
0.77
0.543
D14CP::
5011
0.332
1:384
2514
0.500 cub.9
0335
1:21
0.022
0.200
0./2
1.591
guint)
DT-CIC
4.566
1864
1864 1.465
1.314
887
0.555
01966
0.005
0.355 0.34
0.98 : 0.456
0.113-0-333 0:577
0.001 0.510
1.410
ITT-33:
6300 2 1877
1 908 à 2.136
1/068
1787
1.60 0.317
0:517
n- 187
0.300 0.344
1006 0:337
1.339-0...1 2-711
1667 1177
3.517
ITEAF chown))
1881
1 1119 0.863
1.570
078
0.58 01038
0.137
0.385
0:711
0.33
0.806
5.110
0.90
0.916
finuce)
2.499
8 1199
2563 3903
3005
0.531
0.171
0.150
0.401
0.255
1,315 0.329
0.299
0.127
0.707
1,455
BDT (2-220)
1.92
0.200
1.308 2: 1.854
1,400
1.333
0.318 0.093
0181
2: 2.0013
0.200 0.136
0.196
2 0.031
03392200.154
0.204 0.166
0.63)
0.540 6 0.003
1.225 6 1.526
1/031 8: 1.407
0.142 8 0.000
01995 2: 0.025
0.:20
04029
0,192
2 0/61
(L188 2 0.82)
0.176 8 (4045
0.443
Table 5: The results of BDT on the state-feature distribution matching problem (m = 16). While even simple
auto-encoder regularizers sometimes work well, BDT with longer contexts seems to outperform other strategies and is comparable to CDT or DT (in Table 2). See Appendix E.6 for the full results including = 1.4.
7 CONCLUSION
We provide unified perspective on a wide range of hindsight algorithms, and generalize the problem formulation as hindsight information matching (HIM). Inspired by recent successes in RL as sequence modeling, we propose Generalized Decision Transformer (GDT) which includes DT, Categorical DT, and Bi-directional DT as special cases, and is applicable to any HIM with proper choices of & and aggregator. We show how Categorical DT, a minor modification of DT, enables the first effective offline state-marginal matching algorithm and propose new benchmark tasks for this problem class. We also demonstrate the effectiveness of Bi-directional DT as a one-shot imitation learner, significantly outperforming simple variants based on DT. We hope our proposed HIM and GDT frameworks shed new perspectives on hindsight algorithms and the applicability of sequence modeling to much broader classes of RL problems beyond classic reward-based RI.

---

Published as a conference paper at ICLR 2022
ETHICS STATEMENT
Since this paper mainly focuses on the reinterpretation of hindsight RL algorithms and experiments on existing benchmark datasets, we believe there are no ethical concerns.
REPRODUCIBILITY STATEMENT
We share our codes to ensure the reproductivity. The details of hyperparameters are described in Appendix A and B. The detailed settings of experiments described in the Section 6, Appendix E, and G. We report the results averaged over 20 rollouts every 4 random seed.
REFERENCES
Pieter Abbeel and Andrew Y Ng. Apprenticeship learning via inverse reinforcement learning. In International Conference on Machine learning. 2004.
Marcin Andrychowicz, Filip Wolski, Alex Ray, Jonas Schneider, Rachel Fong, Peter Welinder, Bob McGrew, Josh Tobin, Pieter Abbeel, and Wojeiech Zaremba. Hindsight experience replay. In Advances in neural information processing systems, 2017.
Gabriel Barth-Maron, Matthew W Hoffman, David Budden, Will Dabney, Dan Horgan, Dhruva Tb, Alistair Muldal, Nicolas Heess, and Timothy Lillicrap. Distributed distributional deterministic policy gradients. arXiv preprint arXiv:1804.08617, 2018.
Marc G Bellemare, Will Dabney, and Rémi Munos. A distributional perspective reinforcement learning. In International Conference on Machine Learning, 2017.
Yoshua Bengio, Aaron Courville, and Pascal Vincent. Representation learning: A review and new perspectives. arXiv preprint arXiv: / 253,5338 2012.
Rishi Bommasani, Drew A. Hudson, Ehsan Adeli, Russ Altman, Simran Arora, Sydney von Arx, Michael S. Bernstein, Jeannette Bohg, Antoine Bosselut, Emma Brunskill, Erik Brynjolfsson, Shyamal Buch, Dallas Card, Rodrigo Castellon, Niladri Chatterji, Annie Chen, Kathleen Creel, Jared Quincy Davis, Dora Demszky, Chris Donahue, Moussa Doumbouya, Esin Durmus, Stefano Ermon, John Etchemendy, Kawin Ethayarajh, Li Fei-Fei, Chelsea Finn, Trevor Gale, Lauren Gillespie, Karan Goel, Noah Goodman, Shelby Grossman, Neel Guha, Tatsunori Hashimoto, Peter Henderson, John Hewitt, Daniel E. Ho, Jenny Hong, Kyle Hsu, Jing Huang, Thomas Icard, Saahil Jain, Dan Jurafsky, Pratyusha Kalluri, Siddharth Karamcheti, Geoff Keeling, Fereshte Khani, Omar Khattab, Pang Wei Kohd, Mark Krass, Ranjay Krishna, Rohith Kuditipudi, Ananya Kumar, Faisal Ladhak, Mina Lee, Tony Lee, Jure Leskovec, Isabelle Levent, Xiang Lisa Li, Xuechen Li, Tengyu Ma, Ali Malik, Christopher D. Manning, Suvir Mirchandani, Eric Mitchell, Zancle Munyikwa, Suraj Nair, Avanika Narayan, Deepak Narayanan, Ben Newman, Allen Nie, Juan Carlos Niebles, Hamed Nilforoshan, Julian Nyarko, Giray Ogut, Laurel Orr, Isabel Papadimitriou, Joon Sung Park, Chris Piech, Eva Portelance, Christopher Potts, Aditi Raghunathan, Rob Reich, Hongyu Ren, Frieda Rong, Yusuf Roohani, Camilo Ruiz, Jack Ryan, Christopher Ré, Dorsa Sadigh, Shiori Sagawa, Keshav Santhanam, Andy Shih, Krishnan Srinivasan, Alex Tamkin, Rohan Taori, Armin W. Thomas, Florian Tramèr, Rose E. Wang, William Wang, Bohan Wu, Jiajun Wu, Yuhuai Wu, Sang Michael Xie, Michihiro Yasunaga, Jiaxuan You, Matei Zaharia, Michael Zhang, Tianyi Zhang, Xikun Zhang, Yuhui Zhang, Lucia Zheng, Kaitlyn Zhou, and Percy Liang. On the opportunities and risks of foundation models. arXiv preprint arXiv: 200.:2200.000 100.02288 2021.
Greg Brockman, Vicki Cheung, Ludwig Pettersson, Jonas Schneider, John Schulman, Jie Tang, and Wojciech Zaremba. Openai gym. arXiv preprint 160.vv1600.0155.0 2016.
Tom B. Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhari- wal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert- Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel M. Ziegler, Jeffrey Wu, Clemens Winter, Christopher Hesse, Mark Chen, Eric Sigler, Mateusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCandlish, Alec Radford, Ilya Sutskever, and Dario Amodei. Language models are few-shot learners. arXiv preprint arXiv:2005. 14165, /41655 2020.
10

---

Published as a conference paper at ICLR 2022
Jacob Buckman, Danijar Hafner, George Tucker, Eugene Brevdo, and Honglak Lee. Sample-efficient reinforcement learning with stochastic ensemble value expansion. arXiv preprint arXivv1807.01675. 2018.
Pablo Samuel Castro, Subhodeep Moitra, Carles Gelada, Saurabh Kumar, and Marc G Belle- mare. Dopamine: A research framework for deep reinforcement learning. arXiv preprint arXiv:1812.06110, 2018.
Yevgen Chebotar, Karol Hausman, Yao Lu, Ted Xiao, Dmitry Kalashnikov, Jake Varley, Alex Irpan, Benjamin Eysenbach, Ryan Julian, Chelsea Finn, and Sergey Levine. Actionable models: Unsupervised offline reinforcement learning of robotic skills. arXiv preprint arXiv:2004.07749 2021.
Lili Chen, Kevin Lu, Aravind Rajeswaran, Kimin Lee, Aditya Grover, Michael Laskin, Pieter Abbeel, Aravind Srinivas, and Igor Mordatch. Decision transformer: Reinforcement learning via sequence modeling. arXiv preprint arXiv:2106.01 2021a.
Mark Chen, Jerry Tworek, Heewoo Jun, Qiming Yuan, Henrique Ponde de Oliveira Pinto, Jared Kaplan, Harri Edwards, Yuri Burda, Nicholas Joseph. Greg Brockman, Alex Ray, Raul Puri, Gretchen Krueger. Michael Petrov, Heidy Khlaaf, Girish Sastry, Pamela Mishkin, Brooke Chan, Scott Gray, Nick Ryder, Mikhail Pavlov, Alethea Power, Lukasz Kaiser, Mohammad Bavarian, Clemens Winter, Philippe Tillet, Felipe Petroski Such, Dave Cummings, Matthias Plappert, Fotios Chantzis, Elizabeth Barnes, Ariel Herbert- Voss. William Hebgen Guss, Alex Nichol, Alex Paino, Nikolas Tezak, Jie Tang. Igor Babuschkin, Suchir Balaji, Shantanu Jain, William Saunders, Christopher Hesse, Andrew N. Carr, Jan Leike, Josh Achiam, Vedant Misra, Evan Morikawa, Alec Radford, Matthew Knight, Miles Brundage, Mira Murati, Katie Mayer, Peter Welinder, Bob McGrew, Dario Amodei, Sam McCandlish, Ilya Sutskever, and Wojeiech Zaremba. Evaluating large language models trained on code. arXiv preprini arXiv:2107.03377 2021b.
Jongwook Choi, Archit Sharma, Honglak Lee, Sergey Levine, and Shixiang Shane Gu. Variational empowerment as representation learning for goal-based reinforcement learning. In International Conference on Machine Learning, 2021.
Will Dabney, Mark Rowland, Mare G Bellemare, and Rémi Munos. Distributional reinforcement learning with quantile regression. In Thirty-Second AAAI Conference on Artificial Intelligence. 2018.
Will Dabney, 7eb Kurth-Nelson, Naoshige Uchida, Clara Kwon Starkweather, Demis Hassabis, Rémi Munos, and Matthew Botvinick. A distributional code for value in dopamine-based reinforcement learning. Nature, 577(7792):611-67. 2020.
Sudeep Dasari and Abhinav Gupta. Transformers for one-shot visual imitation. In Conference on Robot Learning, 2020.
Sarah Dean, Horia Mania, Nikolai Matni, Benjamin Recht, and Stephen Tu. On the sample complexity of the linear quadratic regulator. Foundations of Computational Mathematics, 20(4):633-679, 2020.
Jacob Devlin, Ming-Wei Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint arXiv:1810.04805, 2019.
Ron Dorfman, Idan Shenfeld, and Aviv Tamar. Offline meta learning of exploration. arXiv preprint NAXv:20088002598 2021.
Alexey Dosovitskiy. Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al. An image is worth 16x 16 words: Transformers for image recognition at scale. arXiv preprint arXiv:2010.11929, 2020.
Yan Duan, John Schulman, Xi Chen, Peter L Bartlett, Ilya Sutskever, and Pieter Abbeel. RL?: Fast reinforcement learning via slow reinforcement learning. arXiv preprint arXiv: 1. 1.02779, 2016.

---

Published as a conference paper at ICLR 2022
Yan Duan, Marcin Andrychowicz, Bradly C. Stadie, Jonathan Ho, Jonas Schneider, llya Sutskever, Pieter Abbeel, and Wojeiech Zaremba. One-shot imitation learning. In Advances in Neural Information Processing Systems, 2017.
Benjamin Eysenbach, Abhishek Gupta, Julian Ibarz, and Sergey Levine. Diversity is you need: Learning diverse skills without reward function. In International Conference Learning Representations, 2019.
Benjamin Eysenbach, Xinyang Geng, Sergey Levine, and Ruslan Salakhutdinov. Rewriting history with inverse rl: Hindsight inference for policy improvement. In Advances neural information processing systems, 2020.
Benjamin Eysenbach, Ruslan Salakhutdinov, and Sergey Levine. C-learning: Learning achieve goals via recursive classification. In International Conference on Learning Representations, 2021.
Rasool Fakoor, Pratik Chaudhari, Stefano Soatto, and Alexander J. Smola. Meta-q-learning. In International Conference on Learning Representations, 2020.
Chelsea Finn, Paul Christiano, Pieter Abbeel, and Sergey Levine. A connection between generative adversarial networks, inverse reinforcement learning, and energy-basee models. arXiv preprint arXiv:1611.0885 2016.
Chelsea Finn, Tianhe Yu, Tianhao Zhang, Pieter Abbeel, and Sergey Levine. One-shot visual imitation learning via meta-learning, In Conference on Robot Learning, 2017.
Justin Fu, Katie Luo, and Sergey Levine. Learning robust rewards with adverserial inverse reinforce- ment learning. In International Conference on Learning Representations, 2018.
Justin Fu, Aviral Kumar, Ofir Nachum, George Tucker, and Sergey Levine. D4RL: datasets for deep data-driven reinforcement learning. arXiv preprint arXiv:2004.07219, 2020.
Scott Fujimoto and Shixiang Shane Gu. minimalist approach to offline reinforcement learning. arXiv preprint (tl::2006008660 2021.
Scott Fujimoto, Herke van Hoof, and David Meger. Addressing function approximation error in actor-critic methods. In International Conference on Machine Learning, 2018.
Scott Fujimoto, David Meger, and Doina Precup. Off-policy deep reinforcement learning without exploration. In International Conference on Machine Learning, 2019.
Hiroki Furuta, Tadashi Kozuno, Tatsuya Matsushima, Yutaka Matsuo, and Shixiang Shane Gu. Co- adaptation of algorithmic and implementational innovations in inference-based deep reinforcement learning. Advances in neural information processing systems, 202
Hiroki Furuta, Tatsuya Matsushima, Tadashi Kozuno, Yutaka Matsuo, Sergey Levine, Ofir Nachum, and Shixiang Shane Gu. Policy information capacity: Information -heeorti measure for task complexity in deep reinforcement learning. In International Conference on Machine Learning, 2021b.
Seyed Kamyar Seyed Ghasemipour, Shixiang Shane Gu, and Richard Zemel. Smile: Scalable meta inverse reinforcement learning through context-conditional policies. In Advances in Neural Information Processing Systems, 2019.
Seyed Kamyar Seyed Ghasemipour, Richard Zemel, and Shixiang Gu. A divergence minimization perspective on imitation learning methods. In Conference on Robot Learning, 2020.
Seyed Kamyar Seyed Ghasemipour, Dale Schuurmans, and Shixiang Shane Gu. Emaq: Expected-max q-learning operator for simple yet effective offline and online rl. In International Conference on Machine Learning, 2021.
Dibya Ghosh, Abhishek Gupta, Ashwin Reddy, Justin Fu, Coline Manon Devin, Benjamin Eysenbach, and Sergey Levine. Learning to reach goals via iterated supervised learning. In International Conference Learning Representations, 2021.
12

---

Published as a conference paper at ICLR 2022
Adam Gleave and Oliver Habryka. Multi-task maximum entropy inverse reinforcement learning. arXiv preprint 1805. 2018.
Shixiang Gu, Timothy Lillicrap. Ilya Sutskever, and Sergey Levine. Continuous deep q-learning with model-based acceleration. In International Conference on Machine Learning. 2016.
Shixiang Gu, Timothy Lillicrap, Zoubin Ghahramani, Richard Turner, Bernhard Schölkopf, and Sergey Levine. Interpolated policy gradient: Merging on-policy and off -folic gradient estimation for deep reinforcement learning. arXiv preprint 1rXiv:1706.0083 00387, 2017.
Shixiang Shane Gu, Manfred Diaz, Daniel C. Freeman, Hiroki Furuta, Seyed Kamyar Seyed Ghasemipour, Anton Raichuk, Byron David, Erik Frey, Erwin Coumans, and Olivier Bachem. Brax- lines: Fast and interactive toolkit for rl-driven behavior engincering beyond reward maximization. 2021.
arXiv preprint
Jiaming Guo, Rui Zhang, Xishan Zhang, Shaohui Peng, Qi Yi, Zidong Du, Xing Hu, Qi Guo, and Yunji Chen. Hindsight value function for variance reduction in stochastic dynamic environment. arXiv preprint arXiv:2107.12216 2021.
Abhishek Gupta, Benjamin Eysenbach, Chelsea Finn, and Sergey Levine. Unsupervised meta-learning for reinforcement learning. arXiv preprint arXiv: /806 04640, 2018.
Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with a stochastic actor. In International Conference on Machine Learning, 2018
Elad Hazan, Sham M. Kakade, Karan Singh, and Abby Van Soest. Provably efficient maximum entropy exploration. In International Conference Machine Learning. 2019.
Nicolas Heess, Greg Wayne, David Silver, Timothy Lillicrap, Yuval Tassa, and Tom Erez. Learning continuous control policies by stochastic value gradients. arXiv preprint urXiv: 2015.
G. E. Hinton and R. R. Salakhutdinov. Reducing the dimensionality of data with neural networks. Science, (13(6786) 5074 2006.
Jonathan Ho and Stefano Ermon. Generative adversarial imitation learning. In Advances in neural information processing systems, 2016.
Michael Janner, Qiyang Li, and Sergey Levine Reinforcement learning as big sequence modeling problem. arXiv preprint wxXi:2200600203) 2021.
Natasha Jaques, Judy Hanwen Shen, Asma Ghandeharioun, Craig Ferguson, Agata Lapedriza, Noah Jones, Shixiang Shane Gu, and Rosalind Picard. Human-centric dialog training via offline reinforcement learning. arXiv preprint arXiv:2010.05848, 2020.
Nan Jiang, Akshay Krishnamurthy, Alekh Agarwal, John Langford, and Robert E Schapire. Contex- tual decision processes with low bellman rank are pac-learnable. In International Conference on Machine Learning, 2017.
Leslie Pack Kaelbling. Learning achieve goals. In International Joint Conference on Artificial Intelligence, 1993.
Patrick T. Komiske, Eric M. Metodiev, and Jesse Thaler. The Hidden Geometry of Particle Collisions. JHEP, 07:006, 2020. doi: 10. 10.10007/HEP07722000006
Ilya Kostrikov, Kumar Krishna Agrawal, Debidatta Dwibedi, Sergey Levine, and Jonathan Tompson. Diseriminator-actor-riiii Addressing sample inefficiency and reward bias in adversarial imitation learning. In International Conference on Learning Representations, 2019.
Ilya Kostrikov, Ofir Nachum, and Jonathan Tompson. Imitation learning via off-policy distribution matching. In International Conference on Learning Representations, 2020.
Aviral Kumar, Xue Bin Peng, and Sergey Levine. Reward-conditioeed policies. arXiv preprint 1912, 13465, 2019.
13

---

Published as a conference paper at ICLR 2022
Vitaly Kurin, Maximilian Igl, Tim Rocktäschel, Wendelin Boehmer, and Shimon Whiteson. My body is a cage: the role of morphology in graph-based incompatible control. In International Conference on Learning Representations, 2021.
Lisa Lee, Benjamin Eysenbach, Emilio Parisotto, Eric Xing, Sergey Levine, and Ruslan Salakhutdinov. Efficient exploration via state marginal matching. arXiv preprint arXiv:1906.00274. 2020.
Sergey Levine, Aviral Kumar, George Tucker, and Justin Fu. Offline reinforcement Iearning: Tutorial, review, and perspectives open problems. arXiv preprint arXiv:2005.0/64) 2020.
Alexander C. Li, Letrel Pinto, and Pieter Abbeel. Generalized hindsight for reinforcement learning. In Advances in neural information processing systems, 2020.
Lanqing Li, Rui Yang, and Dijun Luo. Focal: Efficient fully-offline meta-reinforremmnt learning via distance metric learning and behavior regularization. In International Conference Learning Representations, 2021.
Yujia Li, Kevin Swersky, and Rich Zemel. Generative moment matching networks. In International Conference Machine Learning, 2015.
Timothy P. Lillicrap. Jonathan J. Hunt, Alexander Pritzel, Nicolas Heess, Tom Erez, Yuval Tassa, David Silver, and Daan Wierstra. Continuous control with deep reinforcement learning. In International Conference Learning Representations, 2016.
Corey Lynch, Mohi Khansari, Ted Xiao, Vikash Kumar, Jonathan Tompson, Sergey Levine, and Pierre Sermanet. Learning latent plans from play. In Conference on Robot Learning, 2019.
Tatsuya Matsushima, Hiroki Furuta, Yutaka Matsuo, Ofir Nachum, and Shixiang Shane Gu. Deployment-efficceen reinforcement learning via model-based offline optimization. In International Conference on Learning Representations, 2021.
Eric Mitchell, Rafael Rafailov, Xue Bin Peng, Sergey Levine, and Chelsea Finn. Offline meta- reinforcement learning with advantage weighting. In International Conference on Machine Learning, 2021.
Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Andrei A Rusu, Joel Veness, Mare G Bellemare, Alex Graves, Martin Riedmiller, Andreas K Fidjeland, Georg Ostrovski, al. Human-level control through deep reinforcement learning. Nature, 2015.
Ofir Nachum, Shixiang Gu, Honglak Lee, and Sergey Levine. Data-Efficient Hierarchical Reinforce- ment Learning. In Advances neural information processing systems, 2018.
Ashvin Nair, Vitchyr Pong, Murtaza Dalal, Shikhar Bahl, Steven Lin, and Sergey Levine. Visual reinforcement learning with imagined goals. arXiv preprint arXiv:1807.04744 2018.
Soroush Nasiriany, Vitchyr H. Pong, Ashvin Nair, Alexander Khazatsky, Glen Berseth, and Sergey Levine. Disco rl: cetrrbbttioocoonitiionnn reinforcement learning for general- purposse policies. In International Conference on Robotics and Automation, 2021.
Emilio Parisotto, Francis Song, Jack Rae, Razvan Pascanu, Caglar Gulcehre, Siddhant Jayakumar, Max Jaderberg. Raphael Lopez Kaufman, Aidan Clark, Seb Noury, al. Stabilizing transformers for reinforcement learning. In International Conference on Machine Learning, 2020.
Lerrel Pinto, Marcin Andrychowicz, Peter Welinder, Wojeiech Zaremba, and Pieter Abbeel. Asym- metric actor critic for image-based robot learning. arXiv preprint 1rXiv:1710.06542, 2017.
Vitchyr Pong, Shixiang Gu, Murtaza Dalal, and Sergey Levine. Temporal difference models: Model- free deep rl for model-based control. International Conference on Learning Representations, 2018.
Vitchyr H. Pong, Ashvin Nair, Laura Smith, Catherine Huang, and Sergey Levine Offline meta- reinforcement learning with online self-supervision. arXiv preprint urXiv:2107.03974, 2021.
Alec Radford, Karthik Narasimhan, Tim Salimans, and Ilya Sutskever. Improving language under- standing by generative pre-training, 2018.
14

---

Published as a conference paper at ICLR 2022
Alec Radford, Jeff Wu, Rewon Child, David Luan, Dario Amodei, and Ilya Sutskever. Language models are unsupervised multitask learners, 2019.
Alec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh, Gabriel Goh, Sandhini Agar- wal, Girish Sastry, Amanda Askell, Pamela Mishkin, Jack Clark, Gretchen Krueger, and llya Sutskever. Learning transferable visual models from natural language supervision. arXiv preprint 10300:2003.00020, 2021.
Kate Rakelly, Aurick Zhou, Deirdre Quillen, Chelsea Finn, and Sergey Levine. Efficient off- ooflpo meta-reinforcemenn learning via probabilistic context variables. In International Conference on Machine Learning, 2019.
Aditya Ramesh, Mikhail Pavlov, Gabriel Goh, Scott Gray, Chelsea Voss, Alec Radford, Mark Chen, and Ilya Sutskever. sero-shot text-to-image generation. arXiv preprint arXiv:2 /2092, 2021.
Zhizhou Ren, Kefan Dong, Yuan Zhou, Qiang Liu, and Jian Peng. Exploration via hindsight goal generation. Advances in neural information processing systems, 2019.
Tim Salimans, Jonathan Ho, Xi Chen, Szymon Sidor, and Ilya Sutskever. Evolution strategies as scalable alternative to reinforcement learning. arXiv preprint arXiv:1703.000 033644 2017.
Archit Sharma, Shixiang Gu, Sergey Levine, Vikash Kumar, and Karol Hausman. Dynamics-awar aware unsupervised discovery of skills. In International conference on learming representations, 2020.
Samarth Sinha, Ajay Mandlekar, and Animesh Garg. S4rl: Surprisingly simple self-supervision for offline reinforcement learning. arXiv preprint arXiv:2103.06326, 2021.
Aravind Srinivas, Michael Laskin, and Pieter Abbeel. Curl: Contrastive unsupervised representations for reinforcement learning. In International Conference on Machine Learning. 2020.
Rupesh Kumar Srivastava, Pranav Shyam, Filipe Mutz, Wojciech Jaškowski, and Jürgen Schmidhuber. Training agents using upside-down reinforcement learning. arXiv preprint arXiv:191..02877, 2019.
Yujin Tang and David Ha. The sensory neuron as a transformer: Permutation-invarian neural networks for reinforcement learning. arXiv preprint arXiv:2109.02869 2021.
Emanuel Todorov, Tom Erez, and Yuval Tassa. Mujoco: A physics engine for model-based control. In International Conference on Intelligent Robots and Systems, 2012.
Stephen Tu and Benjamin Recht. The gap between model-based and model-free methods on the linear quadratic regulator: An asymptotic viewpoint. In Conference on Learning Theory, 2019.
Aaron van den Oord, Yazhe Li, and Oriol Vinyals. Representation learning with contrastive predictive coding. arXiv preprint
2018.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Advances neural information processing systems, 2017.
David Venuto, Flaine Lau, Doina Precup, and Ofir Nachum. Policy gradients incorporating the future. arXiv preprint 10..::21180000990 2021.
Martin J Wainwright and Michael Irwin Jordan. Graphical models, exponential families, and variational inference. Now Publishers Inc, 2008.
Jane X Wang. Zeb Kurth-Nelson, Dhruva Tirumala, Hubert Soyer, Joel Z Leibo, Remi Munos, Charles Blundell, Dharshan Kumaran, and Matt Botvinick. Learning to reinforcement learn. arXiv preprint ar 1rXiv:1611.05700 2016.
Yifan Wu, George Tucker, and Ofir Nachum. Behavior Regularized Offline Reinforcement Learning. arXiv preprint arXiv: /91 1361,, 2019.
Kelvin Xu, Ellis Ratner, Anca Dragan, Sergey Levine, and Chelsea Finn. Learning prior over intent via meta-inverse reinforcement learning. In International Conference Machine Learning, 2019.
15

---

Published as a conference paper at ICLR 2022
Mengjiao Yang and Ofir Nachum. Representation matters:
Offline pretraining for sequential decision making. In International Conference on Machine Learning, 2021.
Lantao Yu, Tianhe Yu, Chelsea Finn, and Stefano Ermon. Meta-inverse reinforcement learning with
probabilistic context variables. In Advances in Neural Information Processing Systems, 2019.
Rui Zhao and Volker Tresp. Energy-basee hindsight experience prioritization. In Conference on Robot Learning, 2018.
Tony Z. Zhao, Jianlan Luo, Oleg Sushkov, Rugile Pevceviciute. Nicolas Heess, Jon Scholz, Stefan Schaal, and Sergey Levine. Offline meta-reinforcement learning for industrial insertion. arXiv preprint arXiv:2 10.04276. 2021.
Brian D Ziebart, Andrew L Maas, J Andrew Bagnell, Anind K Dey, al. Maximum entropy inverse reinforcement learning. In Twenty-Third AAAI Conference on Artificial Intelligence, 2008.
16

---

Published as a conference paper at ICLR 2022
APPENDIX
A HYPER-PARAMETER OF GENERALIZED DECISION TRANSFORMERS
We implement Categorical Decision Transformer and Bi-directional Decision Transformer, built
upon the official codebase released by Chen et al. (2021a) (https: / //g/ttttbbb com/k21/ decision-transforree We follow the most of hyperparameters as they did (Table 6).
Hyperparameter
Value
Number of layers
Number of attention heads
1
Embedding dimension
128
Nonlinearity function
ReLU
Batch size
64
Context length N
Dropout
0.1
Learning rate
le-4
Grad norm clip
0.25
Weight decay
Ic-4
Learning rate decay
Linear warmup for first 100k training steps
Training (Gradient) steps
IM
Number of bins for categorical distribution
31
Encoder size DT-X)
2 layer MLP, (128, 128)
Coefficient of unsupervised loss
0.1
Table 6: List of hyperparameters for DT, CDT and BDT. We refer Chen et al. (2021a).
B DETAILS OF BASELINES
While, to the best of our knowledge, there are no prior work to tackle the offline state-marginal matching problem, we can regard meta or one-shot imitation learning as the methods solving similar problem (see Appendix D for further discussion). In this work, we choose Meta-BC and FOCAL (Li et al., 2021), 1), a metric-based offline meta RL method, as decent baselines.
FOCAL FOCAL is offline meta RL method combining BRAC (Wu et al., 2019) with metric- based approach. It utilizes deterministic context encoder trained with inverse-power distance metric losses and detached from Bellman backup gradients. We follow the hyperparameters in the official implementation (https: gi tuub. com/ LangingLi 1991/99//// (0999///CAA context encoder takes single-step state-action-rewaard tuple /FOCAL- ICLR). Deterministic context for task inference. We parameterize it with (200, 200, 200)-layers MLP. We train deterministic context encoder with 100000 iteration first, and then train BRAC agent with cask-conditioned policy and value functions with 100000 iteration. We use (256, 256, 256)-layen MLPs for policy and value network, and set batch size 256.
Meta-BC We also metric- based MC (Duan et al., 2017; Dasari & Gupta, 2020) strong baseline method for offline multi-task SMM. We adapt deterministic context encoder from FOCAL infer the task. We train Meta-BC in the same way FOCAL just replacing BRAC to BC. The objective is mean-squared error minimization.
Throughout experiments, Meta-BC shows better results than FOCAL, because RL algorithms often struggle to optimize Eq. for distribution matching and FOCAL tries to maximize the task reward, which is not necessarily required for distribution matching problem. In addition, BC often converges faster than offline RL methods.
17

---

Published as a conference paper at ICLR 2022
C DETAILS OF EVALUATION
Throughout the paper, we evaluate both CDT and BDT from a distribution matching perspective by formulating them as offline multi-task SMM or offline multi-task IL problem. However, the current distribution matching research in RL (Lee et al., 2020; Ghasemipour et al., 2020; Gu et al., 2021) has been suffered from the lack of quantitative metrics for evaluation (while standard RL or BC are typically evaluated on the task reward performance). As summarized in Table 7, prior works (Lee et al., 2020; Ghasemipour al., 2020) evaluate the SMM performance by qualitative density visualization using rollout particles.
Although the distance between Mate-marginal and target distribution seems the most intuitive and suitable metric to quantify the performance of distribution matching algorithms, it is often intractable to measure such distance analytically because both state-marginal and target distribution can be non-parametric; cannot access their densities, and only their samples are available. While concurrent work (Gu et al., 2021) tackles this problem by leveraging estimation, introduce single SMM-inspired evaluation for both offline multi-task SMM and offline multi-task II. via estimating Wasserstein-1 distance between empirical categorical distributions. Since the discretization of low-dimensional features, such reward, have success in many RL methods (Bellemare et al., 2017; Dabney et al., 2018; 2020; Furuta et al., 2021b), quantification of distribution matching performance by Wasserstein-1 distance could be reliable evaluations (in addition Wasserstein-1 I distance is symmetric, different from KL distance that is asymmetric). We note that due to the equivalence between inverse RL and SMM methods, the performance of distribution matching methods may be measured via task reward when assuming the accessivity to the expert trajectories (Ho & Ermon, 2016; Fu al., 2018; Kostrikov et al., 2019). However, in our experiments, it is not suitable to evaluate the performance based on task reward since we uses the sub-optimal trajectories (Section 6.1) or multi-task trajectories from different reward functions (Section 6.2).
masple-bseed energy distance
Wasserstein-1 I distance between state-marginal distribution p² (s) and target distribution p" (s) is defined as:
('))))) inf (xx, YIJ,
vEI(p".p")
where I is the set of all possible joint distributions whose marginals are p² and p" respectively (X and Y are random variables). In this work, we empirically estimate Eq. 5 by focusing on the "manually-specified" feature e E F (e.g. reward, xyz-velocities, etc; defined in Section 4) and employing binning and discretization. We discretize the feature space F and obtain B bins (per dimension), whose representatives are {Ộ1, Ø2.... .p). The range of feature space [@min m is pre-defined from the given offline datasets. Then, we get categorical distribution from the target trajectory т: (Ø)
((01,C)). ... (ØB; GB)), and from the rollouts of the policy (using 4 seeds x 20 rollouts); pr (Ø) {(01, ... (0B,ch)); }. where c; and c" are the weights of each bin (i.e. Ee +h and DB B of = 1). We compute the evaluation metric averaging on the test set of the trajectories Drest:
B
B Wiy c
Metric(T, Drest)
Σ
min Wai
c71,
Drest
Diee
i=1 j=1
(6)
B
B Σ
B B
s.t. way 2 0,
Wij S c/ss Σ
Wij S c.
Wiy 1.
J=1
1=1
In practice, we utilize the python package (https: htapp::///ithhnn http//ggttttttteee wtt./com////o//eeco /ttmm////kk/////////emmmmmm
Wasserstein) released by Komiske et al. (2020) to compute it.
Evaluation
Type Qualitative Quantitative
Reference
Density Visualization
Lee al. (2020); Ghasemipour al. (2020)
Energy Distance
Gu et al. (2021)
Task Reward
Quantitative
Ho & Ermon (2016); Fu et al. (2018), etc.
Wasserstein-1 1 Distance
Quantitative
Ours
Table 7: A Review of evaluation for distribution matching algorithms.
18

---

Published as a conference paper at ICLR 2022
D CONNECTION TO META LEARNING
While we solve offlin information matching problems this paper, our formulations (espe- cially Bi-directional DT) and experimental settings are similar offline meta RL (Dorfman et al., 2021; Mitchell et al., 2021; Li et al., 2021) and meta/one-shot imitation learning (Yu et al., 2019; Ghasemipour al., 2019; Finn et al., 2017: Duan al., 2017) (especially metric-based ap- proach (Rakelly et al., 2019; Duan et al., 2016; Fakoor et al., 2020)). We briefly clarify the relevance and difference between them (Table 8).
Offiine Meta RL Recently, some works deal with offline meta RL problem, assuming task distri- bution and offline training with pre-stored transitions collected by the (sub-) optimal or scripted task-conditioned agents (Dorfman et al., 2021; Mitchell et al., 2021; Li et al., 2021; Pong et al., 2021; Zhao et al., 2021). In these works, few test-task trajectories, following the same task distribution but unseen during training phase, are given at test-time (i.e. few-shot), and the agents adapt the given task an online (Pong et al., 2021; Zhao et al., 2021; Dorfman al., 2021) or fully-offline (Mitchell et al., 2021; Li et al., 2021) manner.
Meta Imitation Learning Another related domain is meta imitation learning, also assuming task dis- tribution and offline training with pre-stored expert transitions per tasks (Yu et al., 2019; Ghasemipour Finn et al., 2017; Duan et al., 2017; Xu et al., 2019; Gleave & Habryka, 2018; Dasari & Gupta, 2020). While meta inverse RL methods (Yu et al., 2019; Xu et al., 2019; Gleave & Habryka, 2018) require online samples during test-time adaptation, its off-policy extension and BC-based approach performs offline adaptation with given unseen trajectories (Ghasemipour al., 2019; Finn et al., 2017; Dasari & Gupta, 2020).
et al., 2019;
Our work, contrast, assume no-adaptation at test time and evaluation criteria a distance between the feature distributions, not the task rewards.
Method
Problem
Train
Test
Demo
Pong ct al. (2021)
Offline Meta RL
offline
online
few-shot
Zhao ct al. (2021)
Offline Mcta RL
offline
online
few-shot
Dorfman ct al. (2021)
Offline Meta RL
offline
online
few-shot
Mitchell al. (2021)
Offline Meta RL
offline
offline
few-shot
et (2021)
Offline Meta
offline
offline
few-shot
Yu et al. (2019)
Meta IL
online
online
few-shot
Xu et al. (2019)
Meta IL
online
online
few-shot
Ghasemipour
(2019)
Meta II.
online
offline
few-shot
Finn et al. (2017)
Meta
offline
offline
one-shot
Duan et al. (2017)
Meta II.
offline
(no-adaptation) (no-adaptation)
one-shot
Ours
Offline multi-task SMM/II.
offline
one-shot
Table Review of problem settings among offline RL, meta II, and ours. In offline meta RL, the agents are trained offline and tested with several demonstrations (few-shot), in fully-offfine manner or allowing online data-collection for adaptation. In meta II., IRL-bbased methods train the agents online (Yu et al., 2019; Xu al., 2019; Ghasemipour et al., 2019) while BC-based methods (Finn et al., 2017; Duan et al., 2017) do offline. Similar to settings, Duan et al. (2017) test the agents with single demonstration and no-adaptation process, they evaluate the agent on the task reward, while we do distance between the two distributions.
19

---

Published as a conference paper at ICLR 2022
E DETAILS OF EXPERIMENTS
In this section, we provide the details and supplemental quantitative//ualitttiie results of the experi- ments in Section 6.
Figure 3 visualizes the dataset distributions we use in the experiments. When we sort all the trajectories based on their cumulative rewards, each quality of trajectory (best, middle, worst) shows a different shape of distribution, and the distribution of the whole dataset seems weighted mixture of those.
ger-
PRANO
(a) Reward
(b) X-Velocity
Figure 3: Distributions of the features (reward, and x-velocity) in the D4RL medium-expert datasets.
E. QUANTITATIVE AND QUALITATIVE RESULTS OF REWARD AND STATE-FEATUR MATCHING
We provide the quantitative comaprison of Section 6.1 between Categorical DT and DT in the reward (Table 9) matching problem, computing Wasserstein- 1 distance between the discretized rollout and target feature distributions. Generally, similar to the x-velocity case, CDT shows the better results in offline multi-task SMM compared to original DT. also visualize some of CDT results in Figure 4.
Metaod
Expert
Inlictectah
Median
Total
Expert
Median
Toxl
Expan
miller Mediun
Tital
Ascrage
Cangarioa
0.574 0.213
1.00
0.11 +: 155
11199 o.ss
0.64 no17
0.111 6.077
00855 0.017
3.114
0.105 0.000
0.351
0.652 0
100 1:56
038 à: 138
1337 0.199
8 199 a 6.11
0856 0.015
b676 - 00201
0.541 4 0.45
0-418
BC (oooottee))
8664
0.06
0.119 0007
8.5% 6:34
8.501
0.138
0.682
1,441
Table Quantitative evaluation of reward distribution matching via measuring Wasserstein- distance between the rollout and target distributions. We compare Categorical DT and DT. Since it can capture the multi-modal nature of target distributions, CDT matches the distribution better than the original DT.
(a) Reward
(b) X- X-Velocity
Figure 4: (a) Reward and (b) state-feature (x-velocity) distribution matching in halfcheetah (top), hopper (middle), and walker2d (bottom). The left two examples the distributions from the best trajectories, and right two are the distributions from the middle trajectories in held-out test set. The rollout distributions of CDT (red) match the target distributions (blue) very well in all cases.
20

---

Published as a conference paper at ICLR 2022
E.2 EVALUATION ON TASK PERFORMANCE
While in this paper we focus on the evaluation with the SMM-inspired distribution matching objective, such Wasserstein-1 distance, here provide the evaluation the task rewards. Table 10 shows that DT seems consistently better methods than Categorical DT on the task rewards evaluations, and achieves similar performances the held-out trajectories.
Method
Salfchertah
Medium
Expert
boguer
Medium
Expest
walker2d
Medium
Careorical
5752.57
651.72
312.778
4985.518
1528.442
4055
Held-oul
11146 + 19.070
5237,348
3741 854 7,223
1600 106
0.272
4098
3801
Table 10: Evaluation on the task rewards; conditioning on the held-out trajectories as done in Section 6. We compare the performance between Categorical DT, and DT. DT seems consistently better methods the task rewards evaluations, and achieves similar performances the held-out trajectories (averaged over trajectories).
E.3 2D STATE-FEATURE DISTRIBUTION MATCHING
We also consider two-dimensional state-features (xy-velocities) distribution matching in Ant-v3. Same as ID saatefeeture distribution matching in HalfCheetah, Hopper, and Walker2d-v3 (Section 6.1), we also use cediuum-exper((- dimension separately to reduce the dimension of the categorical distribution that CDT takes as input, while in rest-time we evaluate the performance with Wasserstein- metric on the joint distribution. For DT, we compute the summation of x- and y-velocity each over trajectories and normalize them with the maximum horizon. DT feeds these two scalars as information statistics to match.
datasets from D4RL (Fu et al., 2020). We bin the state-features per
Table 11 reveals that CDT performs better even in the case of two-dimensional state-features dis- tributions, while DT doesn't generalize to expert-quality trajectories. As shown Figure 5, while CDT could cope with the distribution shift between expert and medium target distribution, DT always fits the medium one even the expert trajectory is given as a target. CDT successfully scales to the offiine multi-task SMM problem in the multi-dimensional feature spaces.
Method
ant Medium
Expert
Average
Categorical DT
0.797 + 0.216
0.244 1 0.063
0.521
DT
1.714 0.121
0.260 1 0.067
0.987
Meta-BC
1.295 0.708
0.351 0.205
0.823
FOCAL (Li al., 2021)
1.473 0.892
0.913 + 0.455
1.193
Table 11: Quantitative evaluation of 2D state-feature (xy-velocities) distribution matching, measuring Wasserstein-1 distance. CDT performs better even in the two-dimensional problem.
Target
CDT
DT
2175
0376
0050
0056
0000
0090
3675
027
a590
oc50
0450
0250
are6
ore6
3630
0220
0000
0200
516
512
E12
Medium
Probbbiliy
coa
530
700
620
600
x-Veloiity
x-Veloiity
V-Vlciity
Veloocity
Figure Visualization of 2D state-feature (xy-velocities) distribution matching, binning each dimension separately. Top row shows results from an expert target trajectory, and bottom row from a medium one.
21

---

Published as a conference paper at ICLR 2022
E.4 DETAILS OF SYNTHESIZING UNSEEN BI-MODAL DISTRIBUTION
To construct the dataset, we modified the original reward function in HalfCheetah--3, adding abso- lute z-velocity term (such +np. abs (z_vel)) where the expert cheetah backflips towards -x direction.
We trained SAC agent until convergence (3 million gradient steps), using pytorch implementation released by Furuta et al. (2021a), and then collected 500 trajectories x 1000 time steps. We combined them with 500 trajectories x 1000 time steps from halfcheetah-experrtyy dataset in D4RL, which consists of both backflipping and running forward behaviors.
For the evaluation, we prepared additional 5 trajectories of backflipping and 5 of running forward as uni-modal behaviors (10 test trajectories in total). In addition, we synthesized a bi-modal behavior by dividing each 1000-step trajectories into 500-step sub-trajectories, and concatenating them across different behaviors, which results in the patchworked trajectories of first 500-step running forward and next 500-step backflipping (also 10 trajectories in total).
E.5 DETAILS OF DIVERSE UNSEEN DISTRIBUTION FROM META LEARNING TASK
Following prior meta RL/II works (Rakelly al., 2019; Ghasemipour al., 2019; Pong al., 2021; Li et al., 2021; Fakoor et al., 2020), we modified the reward function for the cheetah to run with specified velocity (such -np. abs (x_vel target_vel) ), and set the horizon 200 steps.
We prepared 31 target x-velocities; taken from [0.0, 3.0], uniformly spaced 0.1 intervals. We also trained the SAC agents until convergence (3 million gradient steps), using pytorch implementation released by Furuta al. (2021a), and collected 250 trajectories 200 time steps each. To simplify the problem than meta learning settings, we held out the 10 trajectories whose x-velocity is (0.5, 1.5, 2.5} as test set, and used the rest as a train data.
E.6 QUANTITATIVE AND QUALITATIVE RESULTS OF ONE-SHOT DISTRIBUTION MATCHING IN FULL STATE
Table 12 and Table 13 show the one-shot reward and state-feature distribution matching results respectively. In addition the embedding size m, also sweep the different context window N = 20, 50, 100 for BDT (m = 16). While even simple auto-encoder regularizer (DT-AE or DT-AE (frozen)) sometimes works well compared no-context BC baselines presented in Table and Table 9, BDT, using (second) anti-causal transformer as an encoder & and its aggregator, seems consistently better than other strategies or Meta-BC baseline for offline multi-task SMM and is comparable to CDT or DT results (also presented in Table 2 and Table 9) with longer context length (N 50). Through the experiment we observe that while there are no clear trends in DT-AE, -CPC or -E2E as the size of context embedding m grows, BDT improves its performance with larger size of embedding. Such intuitive propertics might a good features to design the architectures.
We visualize the qualitative results of BDT (m 16, N 20) in Figure 6.
22

---

Published as a conference paper at ICLR 2022
Medood
Bxnen
lodichertah Mestun
Iotal
Expert
Madzum
bupper
Expers
walker24
Mesiure
Average
DT-CRC (maa))
1.901 0.676
2,35
3935
1.0.11
04000 01037
00
0.09
0.337
0.019
0.139
0.03
0.531
DT-AE (mme)
0.782 0.329
1.720 1,655
1,251
1.301
0633 0.108
0.140 a 0008
0.376
0.233
0.889 0.422
0,265
= -..15
0.77
0433
0.735
DT-CPC (row))
5,887 0.357
1,300 = 1.335
3,023 0: 2.462
0.737 20011
0.238
2 8.000
0447
0.4
1,29
0008
0.14
2 2:000
0.715
0.571
Lot0
(m-15))
1074
1.071
000000
0.37
= 0.125
0.39
DT-CPC
0.304
0.256
0.140
0.712 0.572
1354
1:995
1995
1:30
10ST 10100
11159
02112
11:38
03315
TTN
1XE- (meck, yemr)
4.000
186
2:066 13333
3.348
1.66
0.587 00015
105 1 0.042
0.347
11.44
0801 00711
0750
1 196
0:363
0.57%
D7 AE (nndd join0)
5.563
0.449
1.008
1.365
2.355-
2.458
0.525 + 0.235
0.902
0.433
0.17
0.319
0.215 + 112
n:567
= 0.433
1.585
CRC
0.176
DT (m4e)6
2.166
3.310
1.257
0.655
0.430 0.315
2
0.994
0.41
1.555 0.767
DTCCPC (ow1l, point
4.543
1.779
1,500= 1.433
3.296
1.851
0.577 0.03
0.098 = 0.036
0.341
0.533 + 0.453
0.414 = 0.355
0.663
0,510
1400
DT-E2E (com1)
3:235 2.125
1,00h 1.39
2.123
212
1815 3.0.35
01593 2: 0.205
12218
01690
100 00350
0435
2 0100
01:67 3 0.131
1320
0.551
2.552 0.221
2.888
1.205 2.0990
0.493 0.247
0.949
2835 02946
1.219 2
2107 21 1.91
2733
DIT-F (mall Trooon)
1.355: 1.017
2.506 1.051
1.074
0.563 0100
01511 1/55
0:331
0.340
0.359 2 0006
0.581
1251
DZ CPC (norl, Seazs)
4,100 0.799
2 = -0.78
3.141
1.255
0.582 0/021
0.102
0:041
0.342 0.34
1,222 0.009
0.215 8 0.10
0.545 0.593
LA4
D7 (mse, Inount
0,815
0.183
1,455 = 1.755
1.115
2: 1.353
Dds1 2..00
0.197
2 90061
0449
0.257
1000 (1405
0419
2 021
IL142 2 0.515
0.765
(roms, Sesea))
2:552-
1.158
0.005
0.199 0/662
1.19
03554
0.58
1-444
DT-AE
1.795-
0.25
0.405
117442 (mezt tmons)
1.096
1.164
3.116
0.905
1:345
0436
0199
0:33
0.407 00.28
0.318
0.763
01126
0.00 0.574
1.40
1811 (mE)) 1/300
1.305 11 11:07
17790 1:751
1-88 1255
1031
11185
11100
T 1021
0:001
11117
111T T029
31155
ITNS T 1581
BBT (x=44, N-200-
1.00 0.167
1,288
= 1.580
1.399
1.163
0:494:00.00
0,095 + 0.029
0,355
0.232
0.181
0.053
0,144
0.537
0,298 = 0.397
0.050
BUT (mmit, Ni-ito
1.565
0. 190
1.191
1.830
1.578 1: 1,315
0.11
0.023
01005
1 B017
0.304
2 0.135
0.2H
0815
0.385
1 196
01:00 2: 0.165
0.627
(mmil, X--5))
0,831
0.064
1,204
1.803
LO18 6: 1,200
0.144 0011
0.104
:: 0/236
0.824 0,028
0.199 0.52
0.167
: 0.024
0.183 0044
0.442
ousT Naittn
0.936
0 106
1,300
1,861
100 1,358
0:30
0000
0 8013
0.098
0.07
0,311 0000
0:465 0.991
3:59
Table 12: The results of BDT and DT-X variants on the reward distrihution matching problem (m = 1, 16).
bapper
Flyarr
Walkers
Average
TICAT mma))
1308 TTPI]
X3300 06715
1067
(TTST 1000
01113 10:5
01730 (C155
0.30 11
DT-CRC (moe))
3 0.670
2,277
= 6.349
2.156
0.848
Du631 0.003
0.130
=: 0.037
0,355
0.237
1.777 0.021
0.135 = 0.003
0.600 0.522
L319
DT.AE (mne])
0.789-0033 0.353
1,720
1.714
1.29
6: 1.321
0612 0..108
0.158
= 00066
0,375
0.254
0.884 0,400
0.262
4:033
0.573
4 0433
0.736
1.371
1.347
00008
0.255
(m
2000 1.806
1.055
0.342
0.00
0,321 0:125
0.365
1044 (vente)
6011
0.324
1,405
1.384
3,00
2.514
06:00
0100
B000
0.298
1.281
00022
0.1.6
1.00
0::00
0.572
1.90
17 mal (ami)
1325
1.155
1:90
07000
2305
1:397
10885
01314
11122
1 B055
02877
11235
10710
07811 0 11:67
03317
1-277
CTC (ood, 3000)
4 1.101
2,055
1.055
3:44
1.622
0.05 0.583
0.099
0041
0.341
0.344
0811
0:7711
0.709 -
n 0.519
1.450
(med) jrin
DTC
498
2:210
1989
1.448
0.77
0.341
0.511
0.118
1.571
DTAB (mel5 juie))
0.67
3,300
0.900
5:51
3.764
1.355
0.30 363
0.609
+ D:44
0.053
0.33
101 0.620
0-085 6: 0.413
1.506 0.767
3,000
D7C2 (voel piec)
4.544
1.184
1.866
1,465
3.214
1882
0.515
01032
0,09 + 00025
0.315
0.241
0.00
04454
0.412 > 0.375
ILA80 2: 0.510
1.40
2350
1:20
2133
2:31
1813
U134 0235
0335
1-30)
1.078
9.787
1:097 0.91
0:5
0.541
2.096
0.717
1 0.711
3.517
DT [m=il (ouze)
2238
1:011
2.301 1/855
2553
1.002
0311 0014
0
= 0036
0,35
0.31
1,50
0.145
0:256 =: 01003
11:3 0.500
1223
DT C2C (oodl Sicoaa)
4110= 0.79A
2.705
3.14
1.257
0.578 0222
0.097 = 0000
0.335
0.242
1417
0.000
0.272 =: 109
RL545 = 0.582
1.441
DT (mned, (ruess
1.411 1.783
0.100
0.763
DT-CPC
1.170
0620
0.135
0.25
0.507
0333
0.575
UT-AA (m315, (nsen)
1801
0.582
1.319 0.963
06544 01034
0.157
0.285 (u801
0.128
0-404 0.170
I1/02 0/46k
0.915
1174434 (veals trocns)
3.089
1,399
2.545
0.903
51006
1.141
R631
0190
0.121 à 0:30
11.266
1.51 (4339
0.21% L 12)
(1:9)
0.515
1-415
BBT (rom), Xr-20
1241 01:10
1:1
= 1.270
1.503
1.50
0215
0115
0.00 = B:41
0158
10.10
[00
21135 = 0111
=: 0:100
0.71
1,011
1.175
0.390
0:411
0.405
(mwii8, Xi-20)
1.992 0: 0.301
1
1.854
1,333
0:18 0.093
0.081
0013
0.250
0.136
0.196
0.392
=: 0.194
0.164
0x03
BBT (ww166, NN500
0,840
0.063
1.223=
= 1.825
1001 4: 1.300
0.142 0.00
0.008
= 0/025
0.120 0.029
0.192 0:051
0.163
= 0.02
0.178 2: 0.043
0.443
pa7t coentt, Nest00
0363
0 168
1.306
1.851
11300 1. 1.347
0240 0.544
0.166 1 B032
0.998 2 0.007
0181 0.006
0.533 1 0.14
0:467 0.501
0.304
Table 13: The results of BDT and DT-X variants on the state-feature (x-velocity) distribution matching problem (m 16).
2012
(a) Reward
(b) x- Velociiyy
Figure 6: (a) Reward and (b) state-feature distribution matching by Bi-directional Decision Transformer
(m 16) in halfcheetah (top), hopper (middle), and walker2d (bottom). The left examples are the distributions from the best trajectories, and right two are the distributions from the middle trajectories.
23

---

Published as a conference paper at ICLR 2022
E.7 SHIFTING THE TARGET DISTRIBUTION
To generate the unseen but manageable generalization-tes: trajectories within the support of dataset distribution, make the reward and velocities values of trajectories in the test set shift with constant
offset: bin_size ze X {{- (33.0- -2.0, -1.0, 0.0, +1.0, +2.0, +3.0). Table 14 shows the quantitative comparison between CDT and DT, based on Wasserstein-1 distance between two distributions. CDT successfully handles the distribution shifts (especially in hopper) better than DT. We also provide the state-feature results (Table 15) and qualitative visualizations (Figure 7 and Figure 8), which reveals that CDT can match the rollouts to the shifted target distributions when they are within the support of dataset distributions.
Metted
bopper
Trtal
Tora
Cacqosica
0.101
LA24
=
0.397
1,555
0.041
0,380
0.305
Table 14: Wasserstein-1 distance between shifted reward distribution and the rollout distributions. Categorical DT handles the target distribution shifts and matches the distributions hetter than DT, since CDT aware of distributional information of entire trajectories.
Metlod
halichectah
Total
bopper
Total
Walker3d
Toui
Average
Carposica
0342
2.311
1.21
0.008
0.357
0.247
0052
Table 15: Wasserstein-1 I distance between shifted state-feature (x-velocity) distribution and the rollout distri- butions. Similar to the case of reward, Categorical DT handles the target distribution shifts and matches the distributions better than DTT
watooweaa beg.
Noltheeet ted.
eeraa leot
Tulthheeshh best
Nalchwera beit +57
halkhedah mide
talihesshh madde
nePowetat -1.
Talcheeda do
Telchesahh nadds +11
heltowett ridda +10
bakttedah mith +20
hopper
hopper
Supper
heppor radds-
hepper
bacuee vidde
becper raldo
hepper dd -10
madie
rudde +1.
wwweerribee 2.
wibe 24 bndt
eekercta lea
Waaeeeeee +1
welerrdd ridt +2
waleers hea 450
Ne Bar2o ridde
wellardd rodda
wallor2 midia 0.
Velle Nelee +1.
vellardd ridda +2.0
naloord in col -30
Figure 7: Reward distribution matching in halfcheetah (top two rows; best and middle). hopper (middle two rows; best and middle), and walker2d (bottom two rows; best and middle). We shift the target distribution with (from left right column); bin_size x--3.0 10.0,2...11111. 0.0, +1.0, +2.0, +3.0) (Table 14). Categorical DT (red) can match the rollouts to the shifted target distributions (blue) when the shifted targets are within the support of dataset distributions. For DT (yellow, captioned as Deterministic), we only visualize the delta function at the means of rollouts.
24

---

Published as a conference paper at ICLR 2022
talcherch Teast
talthershh
halfowerh beir
salchaeaa bee
taltherah
huffhwersh
Iefffeettn rokile
Naltoee mere
Neldectin vonde
tekieetinn rakil
Nalcoedat rease
hopper beak 3.
12ppur oet
hopper beat
hopper beat 3
Fosper bet +10
hopper best +23
hogper be0, +30
hopper reledle
hepper roddo 2.5
hopprrrodde 1.0
Toppar
Iepprraicde
maberdd
watoend
5v88 50by
2/1662
Figure State-feature distribution matching, especially x-velocity, halfeheetah (top two rows; best and middle), hopper (middle two rows; best and middle), and walker2d (bottom two rows; best and middle). We shift the target distribution with constant offset (from left to right column); bin_size x(-3.0, 2.0, 11.0, 0.0, +1.0, +2.0, +3.0) (Table 15). For DT (yellow, captioned as Deterministic), we only visualize the delta function the means of rollouts.
25

---

Published as a conference paper at ICLR 2022
E.8 SYNTHESIZING UNREALISTIC TARGET DISTRIBUTION
We synthesize six target distributions manually generating x-velocity samples from Gaussian distribu- tions via python scripts as done in prior works (Chasemipour al., 2020; Gu et al., 2021). Since consider the one-dimensional feature space (x-velocity), we simply specify the mean and standard deviation of Gaussian distributions referring each dataset distribution, as shown in Figure and then generate the 1000 samples per trajectory. While these targets are designed at least within the support of the dataset, we do not consider physical realizabilitt.
HalfCheetah: (14, = (5.0, 1.0), (13.0, 1.0), (9.0, 1.0), (2.5, 1.0), [(5.0, 1.0), (13.0, 1.0)), ((9.0, 1.0), (2.5, 1.0)).
Hopper: (1, (2.5. 1.0), (1.5, 1.0), (3.5, 1.0), (2.5, 0.5), (3..5, 1.0), (1.5, 1.0)), (3..5, 0.5), (1.5, 0.5)].
Walker2d: (1,0) = (3.5, 1.0), (2.5, 1.0), (4.5, 1.0), (1.5, 1.0), ((2.5, 0.5), (4.5, 0.5)). ((1.5, 0.5), (4.5, 0.5)).
For the last two sets, that have different distributional parameters {(L1, 01); (u2, Ø2)}, we sampled from two gaussian distributions 500 samples each, and marge them as trajectory that has multiple modes.
Although they might be unrealistic, Figure 9 implies that CDT tends to match the target distributions, even the cases of bi-modal target distributions. Such generalization to synthesized distribution is an important benefit of citrabutioo-codditionn training. We also quantify the performance of CDT against DT from distributional matching perspective in Table 16.
Inftreettss sytthe and
aytiheedak spstheseed
Felleoston syttresieen
bulcreeton sythesieed
noftreets systhe ant
Topper pptinolesd 1
hopper Dratee
12220 sotthesdedd
bopcer sytteo ool
tesper Hoseeryythoicoa
topper spstholusd
systhesoed
synthe
vakerdd systhetted
systheszedd
Figure 9: State-fcature distribution matching in halfchectah (top), hopper (middle), and walker2d (bottom). We synthesize the each target distribution from Gaussian distributions. While the results worse than for realistic test target distributions in Figure 4, considering that many of these arbitrary synthetic targets could be unrealizable, seeing bi-modal matching results show that indeed CDT has learned generalize something non-trivial.
Method
halfcheetah
hopper
walker2d
Average
Categorical DT
2.059 0.772
0.536 0.225
0.920 1 0.428
1.172
DT
4.256 1 2.220
0.584 1 0.298
0.974 1 0.540
1.938
Table 16: State-feature (x-velocity) distribution matching with synthesized, physically unrealistic target distri- butions generated from scripted Gaussian distributions. We compare Categorical DT and DT. Categorical DT manages deal with such unrealistic target matching.
26

---

Published as a conference paper at ICLR 2022
F DETAILS FOR CATEGORICAL DECISION TRANSFORMER
Categorical Decision Transformer (CDT) takes histograms of categorical distribution (i.e. discrete approximations of feature distributions; B-dim vector) as the inputs of the transformer. Here we describe how to compute the distributions for all timesteps given a trajectory & E [0, ,, T.. For simplicity, we explain the case of one- me-dimensiooall feature (e.g. scalar reward), but for n-dimensional features, we can adopt following procedure per each dimension and arrive at n categorical features. This essentially approximates the full joints with product of independent marginal distribution per dimension, and ensures that number of samples for getting reasonably discretized approximations do not need grow exponentially as the dimension grows.
At first, we discretize the feature space F using B bins (per dimension), and convert the feature o into the one-hot representation di. The range of feature space [0min @max] is pre-defined from the given offline datasets. Ziguse (1), the categorical feature distribution time step I, can be computed recursively following Bellman-like equation:
24hma (t) x कe + 7(1 1lt = 24. (t+ 1),
where 1 is the indicator function. We compute series of 20n in backward manner starting from T. After we obtain the desired information statistics for all trajectories, we feed them to the categorical transformer during training/test time. We describe the python-like pseudocode in Algorithm 1, coloring the changes from the original Decision Transformer (Chen et al., 2021
G DETAILS OF DECISION TRANSFORMER WITH LEARNED Φ
While any unsupervised regularizer for an encoder could be combined into the action MSE loss of Decision Transformer obtain the learned efficiently, observe that even simple objectives, such as auto-encoder and contrastive loss, sometimes perform well. For auto-encoder regularization (DT-AE), we train the MLP encoder (parameterized by v)) and decoder with MSE loss of current state reconstruction:
min E.DD [||s decoder,: (encoder,
Then, use the output of the encoder learned information statistics. In addition, for contrastive loss (DT-CPC), we adopt CURL objective (Srinivas et al., 2020) for state input while adding gaussian perturbation ~ (n 0.0, 0.1) data argumentation (following Sinha et al. (2021)). We train the MLP encoder as query, use its momentum encoder as key:
cxp (encoder, (s) Wencoder, 3 (s + c))
min
log
((..) De'gee exp (encoder, y (s) Wencoder + E!))
where W is learned parameter matrix for the bi-linear inner-product, and J is the weights of the momentum encoder, updated as & <- mah (1 m)), and m 0.95. We treat its trained encoder output as learned information statistics. It remains as future work to combine more advanced, temporary-extended objectives, such as attentive contrastive learning approach proposed in Yang & Nachum (2021). As described in Section 6.3, we consider three strategies to train the encoder for DT-AE and -CPC: training with only unsupervised loss, training with unsupervised and DT's supervised loss jointly (called as "joint"), and pre-training with only unsupervised loss and freezing the weights during DT training (called as "frozen")) We train DT-E2E with only DT's supervised loss as in Algorithm for CDT and BDT.
27

---

Published as a conference paper at ICLR 2022
Algorithm Categoriial//iidireciioo Decision Transformer Pseudocode: Orange and green texts describe additional details on top of the base DT pseudocode from Chen et al. (2021a).
information statistics (histogram, t: states, actions, or timesteps
learned zepresentation)
transformer:
transformer with causal masking (GPT)
embed_s, s, embed_a, embed_z: linear embedding layers
anti Soci_ccasal_tt. second transformer as encoder and aggregator embed t: learned episode positional embedding
pred a: linear action prediction layer compute_stats: function to compute information statistics
main model
def Decisionttraneoommer (z, s, a, t)
# compute embeddings for tokens pos__ _ss_embedding embed_t(t)
s embedding = embed_s (s) + pos _e__ebeddiig
embedding embed_a (a) pos__mmeest pos_embedding
if categorical;
__embedding embed_z (z) pos vos_embedding
elif bi directional:
input state sequence in
order
NOTE: is target state sequence here reversed = fzip(z)
L_ebeedding = embed
tf (zeversed) pos gos_embedding
input embeds = stack (z embedding, embedding, a embedding)
use transformer to get hidden states sidden_states transformer sumputt_mmbeesiintttteeeeee _mmbeddduueeeems
select hidden states for action prediction tokens
hidden unstack thidden_staten, actions
predict action
return pred_a (a hidden)
training loop
train compute stats (train dataset)
dims: (batch size, K, dim)
for 2, (s, a, t) in Strriii z, train dataset):
if unsupervised:
preds DecissonnTramsfrmee (z, loss = mean ( (a (e_prrsss - a) **))
t)
optimizer. zero () loss backward optimizer. step ()
evaluation loop
test z compute stats (test dataset)
t_tests- len (test dataset)
for index in range (n tests):
test cast_trajectory = test_ cast_dataset[[ie [index]
max time time_sttepplp = len (test trajectory)
a, t, done =
[], [1], False
conditioning on the desired information statistics
if categorical:
2 = [test_z [test_ [0]]
elif directional:
z icest_trajectorn [' observations/ [0]]
for range (max_time steps)
action DecissonnTrnssormet (z, t) [-1] new_s, r, done,
env hav.steepee (action)
# append new tokens to sequence if categorical:
z z + [test z [[index] [t+1]] elif bi directional;
z
[test [test_tray]] tesss_traecctor('
observations' obbrevatoon"sssserviiiinnso [t+1]]
s, a, = [new sl, [action], + [len (z)]
z, s, a, z[-N:],
only keep context length
28


Reviews:
[{"review_id": "uNmw9AZkmYe", "reviewer": "Reviewer_5xwG", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Strengths:\n1. Novel work tackling offline state-marginal-matching (SMM) problem. Like other works learning RL skills from offline datasets, this can learn from more signal from fixed data than reward alone. From a learning standpoint, this can aid in learning much more interesting models, especially in the consideration of zero-shot or few-shot learning like GPT or BERT.\n2. By conditioning on the state distribution of a new trajectory, DDT can serve as an imitator at test-time, similar to how language models/GPT can imitate prompts.\n3. Similar to (2) but slightly distinct, this work extends DT to consider a distribution of rewards/returns, which is also novel itself.\n4. Results show fairly strong conditioning performance across the settings considered; in particular I think Section 6.2.2 with an unrealistic target distribution and Section 6.3 with Unsupervised DT are interesting.\n5. Proposed Hindsight Information Matching framework is an interesting way to interpret past work. While I'm not sure about the novelty exactly (i.e. log-probability, squared-error loss are well-known to connect to KL divergence), the work discusses many prior algorithms.\n6. Investigates various representation learning approaches in the context of DTs.\n\nWeaknesses:\n1. While the promise of state-marginal-matching and language modeling is to extract a lot of signal from a dataset, namely trajectory statistics here, much of the experimental analysis focuses on (a) learning from a single feature (as far as I can tell) and (b) x-velocity in particular. On many of these environments, the reward is actually just the x-velocity shifted by some small amount, so it is almost identical to conditioning on the reward. Unsupervised DT lifts the restriction of (a) but is still evaluated on (b). A simple addition would be to evaluate on height; a more substantial addition would be different environments, see (4).\n2. The proposed method seems very similar to one-shot imitation learning (OSIL) [1], where a model is trained to imitate a new trajectory at test time end-to-end, bypassing the statistics step. [2] tackles OSIL with transformers. While DDT/UDT may be more interpretable and/or powerful than OSIL, this is not compared to or discussed in the text.\n3. In the contributions, it is written that the work \"proposed the first benchmark for offline state-marginal matching\". Table 2 is the only quantitative result shown in the paper, however the results seem hard to interpret, and I'm not sure the Wasserstein distance is the right choice of metric here. A proper quantitative metric is important for a benchmark. Adding qualitative comparisons to the main text for better clarity could be helpful. I am also not sure exactly how Deterministic DT works (what is it conditioned on?), but it is probably also a weak baseline. If conditioning on the reward is equivalent to the x-velocity, then Deterministic DT would be a strong baseline. Otherwise, other approaches like OSIL or conditioning on the mean of the target feature would be better.\n4. I think the various representation learning methods investigated for Unsupervised DT are interesting, but they are not really shown or discussed in the main text, so it is hard to think of them as a main investigation of the paper. If the insights can be distilled into a paragraph/table/figure in the main text, I think it would be helpful.\n5. In the introduction, it is written that UDT can \"imitate any test trajectory\". I would strongly recommend rewording this claim, as it is practically impossible to evaluate this claim in continuous settings.\n6. It is not clear to me that the test-time target distributions proposed in Section 6.2 are interesting problems to study and especially to benchmark on. Manipulation tasks, or Atari and other rich environments, could generate more interesting heldout trajectories which are plausible and interesting, but differing in the statistics compared to the training set.\n\nOverall, I am optimistic about this work, and interested in how we can maximally extract training signal from fixed datasets for reinforcement learning. In particular, I think a model which can effectively perform SMM may yield strong finetuning behaviors, although this could be left for future work. Learning to perform SMM offline may also aid in downstream online exploration, which is an important research direction. While I think the experiments and discussion are somewhat weak, I think some relatively simple adjustments (1-5 in weaknesses above) would improve the paper; if they are added, I will change my score to a weak accept. If the authors are interested in continuing this work, I think (6) would make the paper significantly stronger.\n\n=======\n\nMiscellaneous\n\n1. I wasn't fully sure what \"hindsight BC algorithms always work with samples $(\\tau_i, z^{\\tau_i}$\" was supposed to mean; I think hindsight RL methods do this as well? I might be misunderstanding something.\n2. The tables in general have a small font which makes them hard to read, especially Tables 2, 8, and 9 are overloaded. It would be much easier to read with less decimal places or multiplied by 100 or some other choice; an aggregated mean in a single column would also make it easier to interpret. The figure font is also a little small at times.\n3. Under Equation 1 in the definition of $p_t^\\pi(s)$, the $t$ in the product should be redefined to some other $t'$ to not clash with the $t$ representing the time of marginal distribution.\n4. In Equation 5, I think the expectation is missing $s \\sim p_z^\\pi$. (It is included in Equations 4 and 6, for instance)\n5. For the pseudocode in the appendix, I would strongly recommend avoiding the use of red/green due to colorblindness issues.\n\n=======\n\n[1] Yan Duan et al. 2017. \"One-Shot Imitation Learning\"\n\n[2] Sudeep Dasari and Abhinav Gupta 2020. \"Transformers for One-Shot Visual Imitation\"", "summary_of_the_paper": "This work discusses many prior methods under a Hindsight Information Matching (HIM) framework, where the methods can be interpreted as  trying to minimizing the KL divergence between the achieved statistic and some target statistic for a particular choice of statistic. The authors propose to replace the return-to-go conditioning in Decision Transformer with a distribution over some specified state features, which they name Distributional Decision Transformer (DDT). This work also proposes an unsupervised variant, Unsupervised DT (UDT), which does not need certain state features to be specified and can be interpreted as performing offline state-marginal-matching. The authors show in various MuJoCo settings (HalfCheetah, Hopper, Walker2d) that the proposed methods can learn to imitate target information statistics in offline RL and imitation learning settings.", "summary_of_the_review": "The paper is interesting and presents novel methods which may in the future be important ideas leading to new work. Since the results are somewhat weak right now and there are minor clarity concerns, I recommend weak reject.\n\n=====\n\nAfter rebuttal (11/18):\n\nI think the updates made by the authors during the rebuttal period are significant, adding writing clarity, experiments, and the Bidirectional DT. As I wrote below in my response: \"While the scope of experiments is still limited (as also pointed out by Reviewer Ffah), I think the experiments are thorough and the paper contains interesting insights that will be valuable for future investigation, especially given the early state of literature on transformers for RL; I will thus increase my score from 5 to 8.\"", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "j9jlLhaOrel", "reviewer": "Reviewer_cUy5", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "This paper offers three main contributions:\n1. It consolidates a number of existing techniques, including TDMs, Learning from Play, GCSL, and Decision Transformer (DT), into the same umbrella of \"hindsight information matching\" algorithms in the spirit of hindsight relabeling.\n2. It describes a variant of DT, called Categorical DT, that conditions on a desired future state or feature distribution (as opposed to scalar cumulant).\n3. It offers a version of the Categorical DT that includes an unsupervised learning component, with the goal being to learn representations $\\phi$ such that matching the first moment of $\\phi$ is equivalent to matching the distribution of raw features.\n\n**Hindsight information matching**\n\nI did appreciate the exercise of mapping out the space of hindsight algorithms that resulted in Table 1. I am not sure this includes a new insight, since hindsight is commonly acknowledged as a primary motivation for algorithms in this space. For example, from the Decision Transformer paper:\n> Thus, by combining the tools of sequence modeling with **hindsight return information**\n\nand from the Trajectory Transformer paper:\n> ...our method also resembles recently proposed work on goal relabeling [**HER**, GCSL]\nand reward-conditioning [UDRL, RCP] to reinterpret all past experience as useful demonstrations with\nproper contextualization.\n\nand from the GCSL paper:\n> By generating demonstrations using **hindsight relabelling**, we are able to apply goal-conditioned imitation learning primitives...\n\nThat being said, I think it is valuable to aggregate trends like these in a way that makes it easier for researchers to understand how a number of related algorithms fit together, so even if there is nothing technically new from this exposition I do think it is a useful contribution.\n\n**Categorical DT**\n\nI was somewhat less convinced of the other two contributions.\nWhile in principle the ability to match a target _distribution_ should enable substantially more flexible behavior than scalar return-conditioning, this did not seem to be adequately explored or evaluated.\nOne low-level issue is the datasets used for evaluation: all experiments use the `medium-expert` D4RL datasets from what I can tell, which are mixtures of only two policies.\nThis severely limits the types of behaviors that can be extracted from the datasets, and calls into question even some of the more impressive-looking generalization results.\n\nFor example, in Figure 4 we see that the model is able to produce a bimodal x-velocity distribution given a bimodal input (halfcheetah synthesized 6, top right).\nSomewhat suspiciously, however, the modes are in the wrong places.\nThe reason the modes land where they do is because these are the modes of the two policies in the dataset: the medium policy gets an average of 5~6 reward per step, and the expert policy gets an average of ~12 reward per step.\n\nWhile this is possibly the best that one could hope for given the dataset, it sets the bar fairly low.\nSince there are only two behaviors in the dataset, it is only possible to recover different mixtures of two qualitatively different things.\nYou could reduce the entire distributional conditioning to a single parameter that acts as a mixing coefficient between the two policies.\nIn general, this sort of idea would be better evaluated in datasets with higher entropy, to better show the types of behaviors that are recoverable from distribution conditioning that would not be recoverable via standard return conditioning.\n\n**Unsupervised DT**\n\nUnfortunately I could not find quite enough detail about either the method or evaluation here to understand what was going on.\nI think I understand how a representation learning objective could also be used for distribution matching purposes, but aside from that high-level link this addition did come off as a bit tacked-on to the end of another paper.\nMost of the design choices are relegated to the appendix, which is sometimes reasonable given space constraints, but then even the appendix does not provide a self-contained description but mostly references out to other papers, such as:\n> In addition, for contrastive loss, we adopt CURL objective (Srinivas et al., 2020) for state input while removing data augmentation.\n\nMore detail could be useful here -- the data augmentation is important to contrastive methods like CURL, so I am not entirely sure what objective is being used.\n\nThe results are somewhat terse as well:\n> We found that generally SL loss with respect to $\\phi$ could ruin representation learning, as results for SL + UL and SL demonstrate unsuccessful learning , and using contrastive loss is worse than using AE loss (see Figure 15 in Appendix).\n\nIf three of the four variants of Unsupervised DT are meant to serve as either ablations or negative results, could you write more about what we should learn from them?\nI think it could be interesting that the unsupervised-only version works so much better than versions with the supervised loss, but without a more thorough investigation into why this happens it is hard to contextualize.\n\n**Minor**\n1. \"policy that whose future state rollouts\" --> \"policy whose future...\"\n1. Footnote 7: \"prameterized\" --> \"parameterized\"\n1. \" To compute $I_z^\\phi(\\tau_t:T)$ efficiently for all timesteps $t$ given a trajectory $\\tau_1:T$ , we use similar recursive Bellman-like computation inspired by Bellemare et al. (2017).\" Could you write the recursion to make the description more self-contained?\n1. \"CDT don’t match the targets\" --> \"CDT doesn't...\"\n1. \"unsuccessful learning ,\": remove space before comma\n\n**Post-rebuttal update**\n\nThanks to the authors for their responses. I have updated my score based on the conversation below.", "summary_of_the_paper": "This paper presents variants of the Decision Transformer that can condition on desired state or feature distributions instead of scalars.\nAlong the way, the authors also describe an organization of hindsight relabeling methods proposed in the last few years.", "summary_of_the_review": "This paper presents a great and potentially very useful idea, but does not evaluate it in situations where the idea could reasonably be expected to do something more interesting than the original DT.\nBetter motivating the types of behaviors that could be expressed and produced via more flexible conditioning, and seeing that through by testing the method on higher-entropy datasets, would make this a much stronger contribution.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "hfLaNd6YHv", "reviewer": "Reviewer_Ffah", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Strength:\na) unifying framework is developed and prior methods are nicely shown to be special cases\nb) distributional and unsupervised versions of decision transformer are proposed\nc) generalization for 1D distribution shifts is evaluated (i.e., distribution over x-component of velocity is shifted)\n\nWeaknesses:\n1) The paper is somewhat chaotic. It proposes 2 different extensions of decision transformer (DDT and UDT), furthermore introducing two variants of DDT (Categorical (CDT) and Gaussian (GDT)) and two variants of UDT (auto-encoder (AE) and contrastive loss) together with 4 settings (SL+UL, SL, UL, UL-Fix), at the same time it describes offline state-marginal matching and inverse-RL imitation learning and develops a hindsight information matching (HIM) framework. It feels there are enough ideas for several papers. Perhaps because of the scope, the evaluation is rather sketchy. Essentiall, the experiments only show that the basic implementation works, and that generalization w.r.t. to 1D velocity change seems to perform as expected. Focusing on just one variant, e.g., CDT, but providing a more thorough evaluation would be more beneficial and provide a clearer focus for the paper.\n2) Gaussian Decision Transformer is announced but not described. Introduction refers to Sec. 5, which in turn refers to Appendix D, but the latter only covers Categorical Decision Transformer.\n3) The way in which bold numbers are highlighted in tables may be misleading. In the majority of cases, results for DT and DDT lie within 1 standard deviation from each other. It would be more appropriate to either highlight both in bold or not highlight anything.\n4) Only comparison to DT is provided. Comparison to baseline distributional algorithms should be added.", "summary_of_the_paper": "The paper proposes an extension of Decision Transformer (DT) that can work with distributions of features. A number of prior hindsight-based or context-dependent methods are shown to be special cases of a generic scheme called Hindsight Information Matching (HIM), where arbitrary statistics of future trajectories can be used for conditioning. Therefore, the proposed Distributional Decision Transformer (DDT) can be seen as a practical implementation of the HIM algorithm. Experiments on D4RL medium-expert data for HalfCheetah, Hopper, and Walker2D investigate generalization of DDT. Additionally, another extension of DT called Unsupervised Decision Transformer (UDT) is proposed and evaluated that learns the features and rewards in unsupervised manner.", "summary_of_the_review": "The paper is not ready for publication. The unifying perspective presented in the paper is very interesting, as well as the proposed distributional and unsupervised extensions of decision transformer. However, the evaluation is not sufficient. There are too many variations of the algorithm considered and the experiments are not expressive enough to illuminate them. Generalization only studied in a very limited scenario. No comparison to distributional RL is provided. As a result, the paper is trying to cover too many things at once and it is recommended to restrict the scope and provide more thorough analysis.\n\n====\n\nAfter rebuttal\n\nThe authors addressed my concerns with evaluations and baselines. Furthermore, the connection to distributional RL was clarified. The updated paper is 80% new. With the improvements made by the authors, I have no further issues and raise my score to \"8 accept, good paper\".", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Meta-review: The paper describes a framework that unifies several previous lines under hindsight information matching.  Within that framework, the paper also describes variants of the decision transformer (DT) called categorical DT and unsupervised DT.  The rebuttal was quite effective and the reviewers confirmed that their concerns are addressed.  The revised version of the paper is significantly improved and consists of an important contribution that should interested many researchers.  Well done!
Decision: Accept (Spotlight)