Distributional Decision Transformer for Hindsight Information Matching

1 INTRODUCTION . Reinforcement learning ( RL ) suffers from the problem of sample inefficiency , and a 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 dynamics transitions and Bellman equation provide a rich source of supervisory objectives and constraints , many algorithms combined model-free with model-based , and policy-based with value-based in order to achieve maximal sample efficiency , while approximately preserving stable , unbiased policy learning ( Heess et al. , 2015 ; Gu et 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 τt : T to accelerate optimization of a contextual policy π ( at|st , z ) with context z with respect to a parameterized reward function r ( st , at , z ) ( see Section 3 for notations ) . These hindsight algorithms have enabled Q-learning with sparse rewards ( Andrychowicz et al. , 2017 ) , temporally-extended model-based RL with Q-function ( Pong et al. , 2018 ) , mastery of 6-DoF object manipulation in cluttered scenes from human play ( Lynch et al. , 2019 ) , efficient multi-task RL ( Eysenbach et al. , 2020 ; Li et al. , 2020 ) , offline self-supervised discovery of manipulation primitives from pixels ( Chebotar et al. , 2021 ) , and offline RL using return-conditioned supervised learning with transformers ( Chen et al. , 2021a ; Janner et al. , 2021 ) . We derive a generic problem formulation covering all these variants , and observe that this hindsight information matching ( HIM ) framework , with behavioral cloning ( BC ) as 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 as 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 hindsight-inspired algorithms , and Generalized Decision Transformers ( GDT ) as a generalization of DT for RL as sequence modeling to solve any HIM problem ( Figure 1 ) . • Inspired by distribution RL ( Bellemare et al. , 2017 ; Dabney et 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 a 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 ( IL ) , propose a Wasserstein-distance evaluation metric , develop Bi-directional DT ( BDT ) as a fully expressive variant of GDT ( Section 5 ) , and demonstrate BDT ’ s competitive performance at offline multi-task IL . 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. , 2021a ) , the same transition samples can be used to learn with respect to any reward parameters , as long as the reward function is re-computable , i.e . “ relabel ” -able , like goal reaching rewards , and ( 2 ) if policy or Q-functions are 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 to the transitions ’ actions , hindsight methods were applied successfully to enable effective training of goal-based Q-function for sparse rewards ( Andrychowicz et al. , 2017 ) , derive exact connections between Q-learning and classic model-based RL ( Pong et al. , 2018 ) , dataefficient 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 et al. , 2019 ; Zhao & Tresp , 2018 ; Ghosh et al. , 2021 ; Nasiriany et 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 ) , it produced state-of-the-art offline RL and goal-based 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 distribution of returns instead of the averages led to the development of distributional RL algorithms ( Bellemare 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 as discretization and binning , these works focus on optimizing a non-conditional reward-maximizing 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 et al. , 2018 ; Ghasemipour et al. , 2020 ) whose connections to feature-expectation matching have been long discussed ( Abbeel & Ng , 2004 ) . However , those are often exclusively online algorithms even sample-efficient variants ( Kostrikov et al. , 2019 ) , since density-ratio estimations with either discriminative ( Ghasemipour et al. , 2020 ) or generative ( Lee et al. , 2020 ) approach requires on-policy samples , with a rare exception 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 as 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 RL ( 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 ) , 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 statemarginal 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 , a ) , initial state distribution p ( s0 ) , reward function r ( s , a ) , and discount factor γ ∈ ( 0 , 1 ] . In deep RL , a policy that maps the state space to the action space is parameterized by the function approximators , πθ ( a|s ) 1 . The RL objective is given by : LRL ( π ) = 1 1− γ Es∼ρπ ( s ) , a∼π ( ·|s ) [ r ( s , a ) ] ( 1 ) where pπt ( s ) = ∫∫ s0 : t , a0 : t−1 ∏ t p ( st|st−1 , at−1 ) π ( at|st ) and ρπ ( s ) = ( 1 − γ ) ∑ t′ γ t′pπt′ ( st′ = s ) are short-hands for time-aligned and time-aggregated state marginal distributions following policy π . 1For 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 an augmented MDP consisting of past N states , where N is the context window of DT . 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 an alternative problem specification in RL , where instead of stationary-reward maximization , the objective is to find a policy minimizing the divergenceD between its state marginal distribution ρπ ( s ) to a given target distribution p∗ ( s ) 2 : LSMM ( π ) = −D ( ρπ ( s ) , p∗ ( s ) ) ( 2 ) where D is a divergence measure such as Kullback-Leibler ( KL ) 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 a 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 to generate desired behaviors ; and adversarial inverse RL methods ( Ho & Ermon , 2016 ; 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.1 , SMM objectives like Eq.2 no longer depend on task rewards and are only functions of state transition dynamics and target state distribution .

Summary: 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.