Title: Decision Stacks: Flexible Reinforcement Learning via Modular Generative Models
Venue: NeurIPS 2023
Accepted: yes
Keywords: reinforcement learning

generative models

offline RL

sequential decision making
Abstract: Reinforcement learning presents an attractive paradigm to reason about several distinct aspects of sequential decision making, such as specifying complex goals, planning future observations and actions, and critiquing their utilities. However, the combined integration of these capabilities poses competing algorithmic challenges in retaining maximal expressivity while allowing for flexibility in modeling choices for efficient learning and inference. We present Decision Stacks, a generative framework that decomposes goal-conditioned policy agents into 3 generative modules. These modules simulate the temporal evolution of observations, rewards, and actions via independent generative models that can be learned in parallel via teacher forcing. Our framework guarantees both expressivity and flexibility in designing individual modules to account for key factors such as architectural bias, optimization objective and dynamics, transferrability across domains, and inference speed. Our empirical results demonstrate the effectiveness of Decision Stacks for offline policy optimization for several MDP and POMDP environments, outperforming existing methods and enabling flexible generative decision making.

Decision Stacks: Flexible Reinforcement Learning via Modular Generative Models

Abstract

Reinforcement learning presents an attractive paradigm to reason about several distinct aspects of sequential decision making, such as specifying complex goals, planning future observations and actions, and critiquing their utilities. However, the combined integration of these capabilities poses competing algorithmic challenges in retaining maximal expressivity while allowing for flexibility in modeling choices for efficient learning and inference. We present Decision Stacks, a generative framework that decomposes goal-conditioned policy agents into 3 generative modules. These modules simulate the temporal evolution of observations, rewards, and actions via independent generative models that can be learned in parallel via teacher forcing. Our framework guarantees both expressivity and flexibility in designing individual modules to account for key factors such as architectural bias, optimization objective and dynamics, transferrability across domains, and inference speed. Our empirical results demonstrate the effectiveness of Decision Stacks for offline policy optimization for several MDP and POMDP environments, outperforming existing methods and enabling flexible generative decision making. 1

1 Introduction

Modularity is a critical design principle for both software systems and artificial intelligence (AI). It allows for the creation of flexible and maintainable systems by breaking them down into smaller, independent components that can be easily composed and adapted to different contexts. For modern deep learning systems, modules are often defined with respect to their input and output modalities and their task functionalities. For example, Visual ChatGPT [Wu et al., 2023] defines a family of 22+ vision and language foundation models, such as ChatGPT [OpenAI, 2022] (language generation), SAM [Kirillov et al., 2023] (image segmentation), and StableDiffusion [Rombach et al., 2022] (text-to-image generation) for holistic reasoning over text and images. In addition to enabling new compositional applications, modularity offers the promise of interpretability, reusability, and debugging for complex workflows, each of which poses a major challenge for real-world AI deployments.

This paper presents progress towards scalable and flexible reinforcement learning (RL) through the introduction of a new modular probabilistic framework based on deep generative models. Prior work in modular RL focuses on spatiotemporal abstractions that simplify complex goals via hierarchical RL, e.g., [McGovern and Barto, 2001, Andreas et al., 2017, Simpkins and Isbell, 2019, Ahn et al., 2022, Kulkarni et al., 2016, Mendez et al., 2021]. Distinct but complementary to the prior lines of work, our motivating notion of modularity is based on enforcing token-level hierarchies in generative models of trajectories. In the context of RL, trajectories typically consist of a multitude of different tokens of information: goals, observations, rewards, actions. As shown in many recent works [Chen et al., 2021, Janner et al., 2021, 2022, Ajay et al., 2022, Zheng et al., 2022, Reed et al., 2022], we can effectively reduce RL to probabilistic inference [Levine, 2018] via learning deep generative models

over token sequences. However, these frameworks lack any modular hierarchies over the different tokens leading to adhoc choices of generative architectures and objectives, as well as conditional independence assumptions that can be suboptimal for modeling long trajectory sequences.

We introduce Decision Stacks, a family of generative algorithms for goal-conditioned RL featuring a novel modular design. In Decision Stacks, we parameterize a distinct generative model-based module for future observation prediction, reward estimation, and action generation and chain the outputs of each module autoregressively. See Figure 1 for an illustration. While our factorization breaks the canonical time-induced causal ordering of tokens, we emphasize that the relative differences in different token types is significant to necessitate token-level modularity for learning effective policies and planners. Besides semantic differences, the different token types also show structural differences with respect to dimensionalities, domains types (discrete or continuous), modalities (e.g., visual observations, numeric rewards), and information density (e.g., rewards can be sparse, state sequences show relatively high continuity). Instead of modeling the token sequence temporally, parameterizing a distinct module for each token type can better respect these structural differences. In contrast, previous works like Decision Transformer Chen et al. [2021], Trajectory Transformer Janner et al. [2021] and Diffuser Janner et al. [2022] chained states, actions, and, in some cases, rewards within a single temporal model.

In practice, we can train each module in Decision Stacks independently using teacher forcing [Williams and Zipser, 1989], which avoids additional training time as the three modules can be trained in parallel. Decision Stacks shares similarities with many recent works [Janner et al., 2021, Ajay et al., 2022, Janner et al., 2022] that aim to reduce planning to sampling from a generative model. However, our modular design offers additional flexibility and expressivity. Each generative module itself is not restricted to being an autoregressive model and we experiment with modules based on transformers, diffusion models, and novel hybrids. Each generative modeling family makes tradeoffs in architecture, sampling efficiency, and can show varied efficacy for different data modalities. Decision Stacks also brings compositional generalization in scenarios where the modules can be reused across varied tasks or environments. A modular design that easily allows for the use of arbitrary generative models, along with an autoregressive chaining across the modules permits both flexibility and expressivity.

Empirically, we evaluate Decision Stacks on a range of domains in goal-conditioned planning and offline RL benchmarks for both MDPs and POMDPs. We find that the joint effect of modular expressivity and flexible parameterization in our models provides significant improvements over existing offline RL methods. This holds especially in partially observable settings, where Decision Stacks achieves a 15 . 7% performance improvement over the closest baseline, averaged over 9 offline RL setups. We also demonstrate the flexibility of our framework by extensive ablation studies over the choice of generative architectures and inputs for each module.

2 Preliminaries

2.1 Goal conditioned POMDPs

We operate in the formalism of goal-conditioned Partially Observable Markov Decision Processes (POMDP) defined by the tuple M := ( O , S , A , G , P , R , E , γ, p 0 ( s ) , T ) . Respectively, O and S denote the observation space and the underlying state space, which are fully observable in the case of MDPs. A , the action space, is consistent with that of MDPs. In the goal-condition setting, G specifies the task goal distribution which could be e.g., a language instruction or a visual destination state for multi-task policies, or a designed cumulative return for single-task policies. The transition probability function, P : S × A × S → [0 , 1] , describes the transition dynamics. Meanwhile, R : G × S × A ↦→ R defines the rewards that the decision-maker receives after performing action a in state s . The observation emission model E = P ( o | s ) , determines the probability of observing o in state s . The γ , p 0 ( s 0 ) , and T denote the discount factor [Puterman, 2014], initial latent state distribution, and horizon of an episode. In a POMDP context, the observations generated from the underlying state are intrinsically non-Markovian. The goal-conditioned RL objective is to find the optimal policy π ∗ that maximizes the expected cumulative discounted reward over the episode

Observation

Prediction

Reward

Estimation

Action

Generation

P

(

o

1:

T

|

o

0

, G

)

P

(

r

0:

T

|

o

0:

T

, G

)

P

(

a

0:

T

|

o

0:

T

, r

0:

T

, G

)

0.01

0.28

0.41

0.98

Goal

G

Initial Observation

o

0

26.

0.01

0.28

1

0.41

0.98

T = { G,00

r0

ao

01

r1

a1

02

r2

a2

OT

aT

26

→

√

2.2 Offline reinforcement learning

Offline reinforcement learning (RL) is a paradigm for policy optimization where the agent is only given access to a fixed dataset of trajectories and cannot interact with the environment to gather additional samples. Offline RL can be useful in domains where collecting data online is challenging or infeasible such as healthcare [Murphy et al., 2001] and autonomous driving. A major obstacle in offline RL is dealing with distributional shifts. If we naively use Bellman backup for learning the Q-function of a given policy, the update which relies on the actions sampled from policy π can learn inaccurately high values for out-of-distribution actions, leading to instability in the bootstrapping process and causing value overestimation [Kumar et al., 2020].

2.3 Generative models: Autoregressive transformers and Diffusion models

In this work, we are interested in learning the distribution p data ( x | c ) using a dataset D consisting of trajectory samples x and conditioning c . We consider two conditional generative models for parameterizing our agent policies to learn the distribution:

Transformer is a powerful neural net architecture for modeling sequences [Vaswani et al., 2017]. It consists of multiple identical blocks of multi-head self-attention modules and position-wise fullyconnected networks. The vanilla transformer can be modified with a causal self-attention mask to parameterize an autoregressive generative model as in GPT [Radford et al., 2018]. Autoregressive generative models, such as the transformers, factorize the joint distribution p ( x 1 , . . . , x n ) as a product of conditionals, which can be represented as: p ( x ) = ∏ n i =1 p ( x i | x <i ) . This equation shows that the probability of each variable x i depends on the previous variables x 1 , . . . , x i -1 . One advantage of this factorization is that each conditional probability can be trained independently in parallel via teacher forcing [Williams and Zipser, 1989]. In autoregressive generation, sampling is done sequentially, where each variable x i is sampled based on its preceding variables.

Diffusion Models [Sohl-Dickstein et al., 2015, Ho et al., 2020] are latent variable models that consist of a predefined forward noising process q ( x k +1 | x k ) := N ( x k +1 ; √ α k x k , (1 -α k ) I ) that gradually corrupts the data distribution q ( x 0 ) into N (0 , I ) in K steps, and a learnable reverse denoising process p θ ( x k -1 | x k ) := N ( x k -1 | µ θ ( x k , k ) , Σ k ) . For sampling, we first generate a latent sample from the Gaussian prior N (0 , I ) and gradually denoise it using the learned model p θ ( x k -1 | x k ) for

K steps to obtain the data sample x 0 . Diffusion models can be extended for conditional generation using classifier-free guidance [Ho and Salimans, 2022] where a conditional ϵ θ ( x k , c , k ) and an unconditional ϵ θ ( x k , k ) noise model is trained. Conditional data is sampled with the perturbed noise ϵ θ ( x k , k ) + ω ( ϵ θ ( x k , c , k ) -ϵ θ ( x k , k )) , where ω is the guidance strength and c is the conditioning information.

3 Flexible and Modular RL via Decision Stacks

In Figure 2, we consider a directed graphical model for the data generation process in Partially Observable Markov Decision Processes (POMDP) Kaelbling et al. [1998]. The environment encodes the underlying state transitions P ( s t +1 | s t , a t ) , goal-dependent reward function P ( r t | s t , a t , G ) , and the observation emission probability P ( o t | s t ) . Unlike the learning agent which only has access to observations, the behavioral policy used for generating the offline trajectory dataset might also have access to the hidden state information. Such a situation is common in real-world applications, e.g., a human demonstrater may have access to more information about its internal state than what is recorded in video demonstration datasets available for offline learning. Finally, note that a Markov Decision Process (MDP) can be viewed as a special case of a POMDP, where the observation at each timestep, o t , matches the underlying state, s t . To avoid ambiguity, we overload the use of o t to denote both the state and observation in an MDP at time t .

In the context of goal-conditioned decision-making, a finite-horizon trajectory in the offline dataset D is composed of a goal G and a sequence of observations, actions, and reward tokens.

Our primary objective lies in learning a goal-conditioned distribution for P data ( a 0: T , o 1: T , r 0: T | o 0 , G ) conditioned on an arbitrary goal G ∈ G and an initial observation o 0 ∈ O . Leveraging the chain rule of probability, we can factorize this joint distribution into a product of conditional probabilities. For example, Janner et al. [2021] use a time-induced autoregressive factorization:

where τ <t denotes all the tokens in the trajectory before time t . Each conditional factor is parameterized via an autoregressive transformer with shared parameters θ . If the parameterization is sufficiently expressive, any choice of ordering for the variables suffices. However, in practice, we are limited by the size of our offline dataset and the choice of factorization can play a critical role.

In Decision Stacks, we propose to use a modular factorization given as:

Each of the 3 modules (observations, rewards, or actions) focuses on predicting a distinct component of the POMDP and has its own set of parameters ( θ 1 , θ 2 , θ 3 ). Our motivation stems from the fact that in real-world domains, each component is sufficiently distinct from the others in its semantics and representation. Such variances span across a multitude of factors including dimensionalities, domain types (discrete or continuous), modalities (e.g., visual observations, numeric rewards), and information density (e.g., rewards can be sparse, state sequences show relatively high continuity).

Modular Expressivity. In Eq. 3, each module is chained autoregressively with the subsequent modules. This is evident as the output variables of one module are part of the input variables for all the subsequent modules. Under idealized conditions where we can match each module to the data conditional, this autoregressive structure enables maximal expressivity, as autoregressive models derived from the chain rule of probability can in principle learn any data distribution given enough model capacity and training data. Further, our explicit decision to avoid any parameter sharing across modules also permits trivial hardware parallelization and transfer to new environments with shared structure.

Flexible Generative Parameterization. Since each module predicts a sequence of objects, we use any deep generative model for expressive parameterization of each module. Our experiments primarily focus on autoregressive transformers and diffusion models. We also consider hybrid combinations, as they are easy to execute within our framework and can avoid scenarios where individual model families suffer, e.g., diffusion models lag behind transformer for discrete data; whereas transformers are generally poor for modeling continuous signals such as image observations. In real-world environments, many of these challenges could simultaneously occur such as agents executing discrete actions given continuous observations. Finally, each module is conditioned on a goal. For training a multi-task agent, the goal can be specified flexibly as spatial coordinates, a visual image, a language instruction, etc. For single-task agents, we specify the goal as the trajectory returns during training and desired expert-valued return during testing, following prior works in return-conditioned offline RL [Chen et al., 2021, Emmons et al., 2021].

Learning and Inference. Given an offline dataset, each module can be trained in parallel using teacher forcing [Williams and Zipser, 1989]. At test-time, our framework naturally induces a planner, as in order to predict an action at time t , we also need to predict the future observations and rewards. We can execute either an open-loop or closed-loop plan. Open-loop plans are computationally efficient as they predict all future observations and rewards at once, and execute the entire sequence of actions. In contrast, a closed-loop plan is likely to be more accurate as it updates the inputs to the modules based on the environment outputs at each time-step. Using a closed-loop plan, we can sample the action at time t as follows:

The hat symbol (^) indicates predicted observations, rewards, and actions, while its absence denotes observations, rewards, and actions recorded from the environment and the agent in the previous past timesteps. For closed-loop planning, Eqs. 4, 5, 6 require us to condition the joint observation, reward and action distributions on the past trajectory tokens. For a module that is parameterized autoregressively, this is trivial as we can simply choose a time-induced ordering and multiply the conditionals for the current and future timesteps. For example, if the observation module is an autoregressive transfer, then we can obtain the sampling distribution in Eq. 4 as: P θ 1 ( o t +1: T | o 0: t , G ) = ∏ T i = t +1 P θ 1 ( o i | o <i , G ) . For a diffusion model, this task is equivalent to inpainting and can be done by fixing the environment observations until time t at each step of the denoising process [Janner et al., 2022].

Distinction with Key Prior Works. We will include a more detailed discussion of broad prior works in §5 but discuss and contrast some key baselines here. While the use of generative models for goal-conditioned offline RL is not new, there are key differences between Decision Stacks and recent prior works. First, we choose a planning approach unlike other model-free works, such as Decision Transformers [Chen et al., 2021] and diffusion-based extensions [Wang et al., 2022]. Second, there exist model-based approaches but make different design choices; Trajectory Transformer [Janner et al., 2021] uses a time-induced causal factorization parameterized by a single autoregressive transformer, Diffuser [Janner et al., 2022] uses diffusion models over stacked state and action pairs, Decision Diffuser [Ajay et al., 2022] uses diffusion models for future state prediction and an MLP-based inverse dynamics model to extract actions. Unlike these works, we propose a modular structure that is maximally expressive as it additionally models reward information and does not make any conditional independence assumption for the state, reward and action modules. As our experiments demonstrate, the modular expressivity and architectural flexibility in Decision Stacks are especially critical for goal-conditioned planning and dealing with partial observability.

4 Experiments

Our experiments aim to answer the following questions:

§4.1 How does Decision Stacks perform for long-horizon multi-task planning problems? §4.2 How does Decision Stacks compare with other offline RL methods in MDP environments? §4.3 How does Decision Stacks compare with other offline RL methods in POMDP environments? §4.4 How does the architectural feasibility for each module affect downstream performance? How does the modularity enable compositional generalization? How important is the role of reward modeling for Decision Stacks?

For §4.1, §4.2, and §4.3, we experiment with D4RL environments and parameterize Decision Stacks with a diffusion-based observation model, an autoregressive transformer-based reward model, and an autoregressive transformer-based action model. Finally, in §4.4, we will ablate the full spectrum of architecture design choices for each module.

4.1 Long-Horizon Goal-Conditioned Environments

We first test for the planning capabilities of Decision Stacks on the Maze2D task from the D4RL [Fu et al., 2020] benchmark. This is a challenging environment requiring an agent to generate a plan from a start location to a goal location. The demonstrations contain a sparse reward signal of +1 only when the agent reaches close to the goal. Following Janner et al. [2022], we consider 2 settings. In the Single Goal setting, the goal coordinates are fixed, and in the Multi Goal setting, the goals are randomized at test-time. We compare against classic trajectory optimization techniques that have knowledge of the environment dynamics (MPPI [Williams et al., 2015], extensions of model-free RL baselines (CQL [Kumar et al., 2020] and IQL [Kostrikov et al., 2021]), and the two most closely related works in generative planning based on diffusion models: Diffuser [Janner et al., 2022] and Decision Diffuser (DD) [Ajay et al., 2022].

We show our results in Table 1 for different goal types and maze grids. While Janner et al. [2022] previously demonstrated remarkable ability in generating long-horizon plans using Diffuser, their trajectory plans were executed by a handcoded controller. However, we experimentally found that Diffuser and DD's own generated actions fail to perfectly align with their generated plans, as shown in the example rollouts in Figure 3. We hypothesize this could stem from the lack of modularity in Diffuser affecting the generation fidelity, or the lack of expressivity in using an MLP-based inverse dynamics model in DD which limits the context length required for long-horizon planning. In contrast, we find that DS generates robust trajectory plans and matching action sequences with significant improvements over baselines.

4.2 Offline Reinforcement Learning Performance in MDPs

Next, we examine the performance of Decision Stacks in offline RL tasks across various highdimensional locomotion environments from the D4RL offline benchmark suite [Fu et al., 2020] in Table 2. We compare Decision Stacks (DS) with other offline RL algorithms including imitation learning via Behavior Cloning (BC), value-based approaches like IQL [Kostrikov et al., 2021] and CQL [Kumar et al., 2020], model-based algorithm MOReL [Kidambi et al., 2020], transformer-based

Diffuser

DD

DS

Diffuser with

Goal

DS / DD with

.5

generative models such as Decision Transformer (DT) [Chen et al., 2021] and Trajectory Transformer (TT) [Janner et al., 2021], and diffusion-based generative models Diffuser [Janner et al., 2022] and Decision Diffuser (DD) [Ajay et al., 2022]. In our evaluation, we also included our reproduced scores for DD. DD uses the same architecture for observation prediction as Decision Stacks and is hence, the closest baseline. However, we found its performance to be sensitive to return conditioning and in spite of an extensive search for hyperparameters and communication with the authors, our reproduced numbers are slightly lower. We provide more details in the Appendix. For a fair comparison, we used the same set of hyperparameters that give the best performance for the DD baseline.

We show results averaged over 15 planning seeds and normalize the scores such that a value of 100 represents an expert policy, following standard convention [Fu et al., 2020]. Decision Stacks outperforms or is competitive with the other baselines on 6/9 environments and is among the highest in terms of aggregate scores. These results suggest that even in environments where we can make appropriate conditional independence assumptions using the MDP framework, the expressivity in the various modules of Decision Stacks is helpful for test-time generalization.

4.3 Offline Reinforcement Learning Performance in POMDPs

Next we consider the POMDP setting where the logged observations are incomplete representations of the underlying states. To generate the POMDPs datasets, we exclude the two velocity dimensions from the full state representation, which consists of both positions and velocities. This simulates a lack of relevant sensors. A sensitivity analysis on dimension occlusions further strengthens our results, as shown in Appendix Table 7. DS continues to outperform other baselines for each environment from the D4RL locomotion datasets. We report our results in Table 3 and compare against other generative baselines. Decision Stacks (DS), consistently achieves competitive or superior results compared to the other algorithms, including BC, DT, TT, and DD. Notably, DS outperforms other methods in

most environments and attains the highest average score of 74 . 3 , which reflects a 15 . 7% performance improvement over the next best-performing approach Diffuser. This highlights the effectiveness of our approach in handling POMDP tasks by more expressively modeling the dependencies among observations, actions, and rewards.

4.4 Architectural Flexibility and Compositional Generalization

Decision Stacks distinctly separates the prediction of observations, rewards, and actions employing three distinct models that can be trained independently using teacher forcing. In this section, we explore the additional flexibility offered by different architecture choices for each module. For observation, reward, and action prediction, we consider diffusion models and Transformer-based autoregressive models. For reward and action models, we additionally consider MLPs that are restricted in their window and only look at the immediate state information to make a decision. The

results shown in Table 4 display a combination of 2x3x3 policy agents for the Hopper-medium v2 POMDP environment. Since we adopt a modular structure, we can compose the different modules efficiently and hence, we only needed to train 2 (state) + 3 (reward) + 3 (action) models. In Table 4, we find that the performance of pure transformer- or diffusion-based Decision Stacks gives reasonable performance (transformers: 53.0, diffusion: 56.9) but these pure combinations can be slightly outperformed by hybrids, e.g., the best achieving entry (58.2) in Table 4 uses a diffusion-based obsevation model, a transformer-based reward model and a diffusion-based action model. MLPs generally are outperformed by generative architectures, especially when used for modeling actions.

In addition to architectural flexibility, DS's modularity also supports compositional transfer across tasks. The lack of parameter sharing across modules not only enables hardware parallelization but also enables the transfer of modules to new environments. As shown in table 6, we created modified versions of the Maze2d-Umaze-v2 environment, altering action and reward spaces while keeping the observation model consistent. Both the reward and observation models were kept consistent during transfers between different action spaces. A single observation model was shared across six environments, and each of the two reward models was utilized across three action spaces. Our findings show that DS promotes efficient learning and compositional generalization, emphasizing its modular efficiency in scenarios with reusable components across varied tasks or environments.

Table 6: Experiments on Compositional Generalization. In the Maze2d-Umaze-v2 environment, we analyze both dense and sparse rewards. For the action space, we consider three variations of the force on a 2D ball: 1) Unscaled Force: the original force applied to the ball in the 2D maze environment; 2) Reflected Force: the original force mirrored across the x-axis; 3) Rotated Force: the original force rotated 30 degrees counterclockwise, altering the direction of the applied force. Each row shares the same reward model and observation model across the three different action spaces. The results underscore the modular efficiency of DS in scenarios where reusable components exist across different tasks or environments. The standard error is reported, and the results are averaged over 15 random seeds.

Furthermore, we compare the best reward modeling architectures with alternatives that do not consider rewards. This is standard practice for Diffuser [Janner et al., 2022] and Decision Diffuser (DD) [Ajay et al., 2022]. For example, DD predict the action at time t only based on the current observation and next observation, P ( a t | o t , o t +1 ) , parameterized via an MLP. As delineated in Table 5, the inclusion of reward models significantly boosts performance in the dense reward POMDP environment. We include additional analysis and discussion in the Appendix D.

5 Related works

Offline Reinforcement Learning is a paradigm that for learning RL policies directly from previously logged interactions of a behavioral policy. The key challenge is that any surrogate models trained on an offline dataset do not generalize well outside the dataset. Various strategies have been proposed to mitigate the challenges due to distribution shifts by constraining the learned policy to be conservative

and closely aligned with the behavior policy. These include learning a value function that strictly serves as a lower bound for the true value function in CQL [Kumar et al., 2020], techniques focus on uncertainty estimation such as Kumar et al. [2019], and policy regularization methods [Wu et al., 2019, Ghasemipour et al., 2021, Kumar et al., 2019, Fujimoto and Gu, 2021, Fujimoto et al., 2019]. Model-based methods like MORel [Kidambi et al., 2020] and ReBeL [Lee et al., 2021] add pessimism into the dynamics models. In the context of partially observed settings, Rafailov et al. [2021] extends model-based offline RL algorithms by incorporating a latent-state dynamics model for high-dimensional visual observation spaces, effectively representing uncertainty in the latent space and Zheng et al. [2023] derive algorithms for offline RL in settings where trajectories might have missing actions. Our work takes the RL as inference perspective [Levine, 2018] and employs the tools of probabilistic reasoning and neural networks for training RL agents.

Generative models for offline RL. Over the past few years, the RL community has seen a growing interest in employing generative models for context-conditioned sequence generation by framing the decision-making problem as a generative sequence prediction problem. Here, we expand on our discussion from §3 with additional context and references. Decision Transformer [Chen et al., 2021] and Trajectory Transformer [Janner et al., 2021] concurrently proposed the use of autoregressive transformer-based models [Radford et al., 2018] for offline RL in model-based and model-free setups. Online decision transformer [Zheng et al., 2022] further finetunes the offline pretrained policies in online environments through a sequence-level exploration strategy. GATO [Reed et al., 2022, Lee et al., 2022] and PEDA [Zhu et al., 2023] scale these models to multi-task and multi-objective settings. MaskDP Liu et al. [2022] shows that other self-supervised objectives such as masking can also enable efficient offline RL, especially in goal-conditioned settings. These advancements have also been applied to other paradigms in sequential decision making such as black-box optimization and experimental design [Nguyen and Grover, 2022, Krishnamoorthy et al., 2023a,b]. Recent works have shown that changing the generative model from transformer to a diffuser with guidance can improve performance in certain environments and also permit planning for model-based extensions [Janner et al., 2021, Ajay et al., 2022, Wang et al., 2022, Chen et al., 2022]. Dai et al. [2023] considers an extension where the state model is a pretrained text2image model and actions are extracted from consecutive image frames. As discussed in §3, these works make specific design choices that do not guarantee the modularity, flexibility, and expressivity ensured by our framework.

6 Conclusion

We proposed Decision Stacks, a modular approach for learning goal-conditioned policies using offline datasets. Decision Stacks comprises of 3 modules tasked with the prediction of observations, rewards, and actions respectively. In doing so, we strive for the twin benefits of expressivity through autoregressive conditioning across the modules and flexibility in generative design within any individual module. We showed its empirical utility across a range of offline RL evaluations for both MDP and POMDP environments, as well as long-horizon planning problems. In all these settings, Decision Stacks matches or significantly outperforms competing approaches while also offering significant flexibility in the choice of generative architectures and training algorithms.

Limitations and Future Work. Our experiments are limited to state-based environments and extending Decision Stacks to image-based environments is a promising direction for future work especially in light of the gains we observed for POMDP environments. We are also interested in exploring the benefits of a modular design for pretraining and transfer of modules across similar environments and testing their generalization abilities. Finally, online finetuning of Decision Stacks using techniques similar to Zheng et al. [2022] is also an exciting direction of future work.

Acknowledgments

This research is supported by a Meta Research Award.

References

Ahn, A. Brohan, N. Brown, Y. Chebotar, O. Cortes, B. David, C. Finn, C. Fu, K. Gopalakrishnan, K. Hausman, A. Herzog, D. Ho, J. Hsu, J. Ibarz, B. Ichter, A. Irpan, E. Jang, R. J. Ruano, K. Jeffrey, S. Jesmonth, N. Joshi, R. Julian, D. Kalashnikov, Y. Kuang, K.-H. Lee, S. Levine, Y. Lu, L. Luu, C. Parada, P. Pastor, J. Quiambao, K. Rao, J. Rettinghouse, D. Reyes, P. Sermanet, N. Sievers, C. Tan, A. Toshev, V. Vanhoucke, F. Xia, T. Xiao, P. Xu, S. Xu, M. Yan, and A. Zeng. Do as i can and not as i say: Grounding language in robotic affordances. In arXiv preprint arXiv:2204.01691 , 2022.

Ajay, Y. Du, A. Gupta, J. Tenenbaum, T. Jaakkola, and P. Agrawal. Is conditional generative modeling all you need for decision-making? arXiv preprint arXiv:2211.15657 , 2022.

Andreas, D. Klein, and S. Levine. Modular multitask reinforcement learning with policy sketches. In International Conference on Machine Learning , pages 166-175. PMLR, 2017.

Chen, C. Lu, C. Ying, H. Su, and J. Zhu. Offline reinforcement learning via high-fidelity generative behavior modeling. arXiv preprint arXiv:2209.14548 , 2022.

Chen, K. Lu, A. Rajeswaran, K. Lee, A. Grover, M. Laskin, P. Abbeel, A. Srinivas, and I. Mordatch. Decision transformer: Reinforcement learning via sequence modeling. Advances in neural information processing systems , 34:15084-15097, 2021.

Dai, M. Yang, B. Dai, H. Dai, O. Nachum, J. Tenenbaum, D. Schuurmans, and P. Abbeel. Learning universal policies via text-guided video generation. arXiv preprint arXiv:2302.00111 , 2023.

Emmons, B. Eysenbach, I. Kostrikov, and S. Levine. Rvs: What is essential for offline rl via supervised learning? arXiv preprint arXiv:2112.10751 , 2021.

Fu, A. Kumar, O. Nachum, G. Tucker, and S. Levine. D4rl: Datasets for deep data-driven reinforcement learning. arXiv preprint arXiv:2004.07219 , 2020.

Fujimoto and S. S. Gu. A minimalist approach to offline reinforcement learning. Advances in neural information processing systems , 34:20132-20145, 2021.

Fujimoto, D. Meger, and D. Precup. Off-policy deep reinforcement learning without exploration. In International conference on machine learning , pages 2052-2062. PMLR, 2019.

K. S. Ghasemipour, D. Schuurmans, and S. S. Gu. Emaq: Expected-max q-learning operator for simple yet effective offline and online rl. In International Conference on Machine Learning , pages 3682-3691. PMLR, 2021.

Ho and T. Salimans. Classifier-free diffusion guidance. arXiv preprint arXiv:2207.12598 , 2022.

Ho, A. Jain, and P. Abbeel. Denoising diffusion probabilistic models. Advances in Neural Information Processing Systems , 33:6840-6851, 2020.

Janner, Q. Li, and S. Levine. Offline reinforcement learning as one big sequence modeling problem. Advances in neural information processing systems , 34:1273-1286, 2021.

Janner, Y. Du, J. B. Tenenbaum, and S. Levine. Planning with diffusion for flexible behavior synthesis. arXiv preprint arXiv:2205.09991 , 2022.

P. Kaelbling, M. L. Littman, and A. R. Cassandra. Planning and acting in partially observable stochastic domains. Artificial intelligence , 101(1-2):99-134, 1998.

Kidambi, A. Rajeswaran, P. Netrapalli, and T. Joachims. Morel: Model-based offline reinforcement learning. Advances in neural information processing systems , 33:21810-21823, 2020.

P. Kingma and J. Ba. Adam: A method for stochastic optimization. In International Conference on Learning Representations (ICLR) , 2015.

Kirillov, E. Mintun, N. Ravi, H. Mao, C. Rolland, L. Gustafson, T. Xiao, S. Whitehead, A. C. Berg, W.-Y. Lo, P. Dollár, and R. Girshick. Segment anything. arXiv:2304.02643 , 2023.

Kostrikov, A. Nair, and S. Levine. Offline reinforcement learning with implicit q-learning. arXiv preprint arXiv:2110.06169 , 2021.

Krishnamoorthy, S. M. Mashkaria, and A. Grover. Generative pretraining for black-box optimization. In International Conference on Machine Learning (ICML) , 2023a.

Krishnamoorthy, S. M. Mashkaria, and A. Grover. Diffusion models for black-box optimization. In International Conference on Machine Learning (ICML) , 2023b.

D. Kulkarni, K. Narasimhan, A. Saeedi, and J. Tenenbaum. Hierarchical deep reinforcement learning: Integrating temporal abstraction and intrinsic motivation. Advances in neural information processing systems , 29, 2016.

Kumar, J. Fu, M. Soh, G. Tucker, and S. Levine. Stabilizing off-policy q-learning via bootstrapping error reduction. Advances in Neural Information Processing Systems , 32, 2019.

Kumar, A. Zhou, G. Tucker, and S. Levine. Conservative q-learning for offline reinforcement learning. Advances in Neural Information Processing Systems , 33:1179-1191, 2020.

B.-J. Lee, J. Lee, and K.-E. Kim. Representation balancing offline model-based reinforcement learning. In International Conference on Learning Representations , 2021.

K.-H. Lee, O. Nachum, M. S. Yang, L. Lee, D. Freeman, S. Guadarrama, I. Fischer, W. Xu, E. Jang, H. Michalewski, et al. Multi-game decision transformers. Advances in Neural Information Processing Systems , 35:27921-27936, 2022.

Levine. Reinforcement learning and control as probabilistic inference: Tutorial and review. arXiv preprint arXiv:1805.00909 , 2018.

Liu, H. Liu, A. Grover, and P. Abbeel. Masked autoencoding for scalable and generalizable decision making. In Advances in Neural Information Processing Systems (NeurIPS) , 2022.

McGovern and A. G. Barto. Automatic discovery of subgoals in reinforcement learning using diverse density. In International Conference on Machine Learning (ICML) , 2001.

A. Mendez, H. van Seijen, and E. EATON. Modular lifelong reinforcement learning via neural composition. In International Conference on Learning Representations , 2021.

Misra. Mish: A self regularized non-monotonic activation function. arXiv preprint arXiv:1908.08681 , 2019.

A. Murphy, M. J. van der Laan, J. M. Robins, and C. P. P. R. Group. Marginal mean models for dynamic regimes. Journal of the American Statistical Association , 96(456):1410-1423, 2001.

Nguyen and A. Grover. Transformer neural processes: Uncertainty-aware meta learning via sequence modeling. In International Conference on Machine Learning (ICML) , 2022.

Nguyen, Q. Zheng, and A. Grover. Reliable conditioning of behavioral cloning for offline reinforcement learning. arXiv preprint arXiv:2210.05158 , 2022.

OpenAI. Openai: Introducing chatgpt, 2022. URL https://openai.com/blog/chatgpt .

L. Puterman. Markov decision processes: discrete stochastic dynamic programming . John Wiley &Sons, 2014.

Radford, K. Narasimhan, T. Salimans, I. Sutskever, et al. Improving language understanding by generative pre-training. In preprint, OpenAI , 2018.

Rafailov, T. Yu, A. Rajeswaran, and C. Finn. Offline reinforcement learning from images with latent space models. In Learning for Dynamics and Control , pages 1154-1168. PMLR, 2021.

Reed, K. Zolna, E. Parisotto, S. G. Colmenarejo, A. Novikov, G. Barth-Maron, M. Gimenez, Y. Sulsky, J. Kay, J. T. Springenberg, T. Eccles, J. Bruce, A. Razavi, A. Edwards, N. Heess, Y . Chen, R. Hadsell, O. Vinyals, M. Bordbar, and N. de Freitas. A generalist agent, 2022.

Rombach, A. Blattmann, D. Lorenz, P. Esser, and B. Ommer. High-resolution image synthesis with latent diffusion models. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR) , pages 10684-10695, June 2022.

Simpkins and C. Isbell. Composable modular reinforcement learning. In Proceedings of the AAAI conference on artificial intelligence , volume 33, pages 4975-4982, 2019.

Sohl-Dickstein, E. Weiss, N. Maheswaranathan, and S. Ganguli. Deep unsupervised learning using nonequilibrium thermodynamics. In International Conference on Machine Learning , pages 2256-2265. PMLR, 2015.

Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A. N. Gomez, Ł. Kaiser, and I. Polosukhin. Attention is all you need. Advances in neural information processing systems , 30, 2017.

Wang, J. J. Hunt, and M. Zhou. Diffusion policies as an expressive policy class for offline reinforcement learning. arXiv preprint arXiv:2208.06193 , 2022.

Williams, A. Aldrich, and E. Theodorou. Model predictive path integral control using covariance variable importance sampling. arXiv preprint arXiv:1509.01149 , 2015.

J. Williams and D. Zipser. A learning algorithm for continually running fully recurrent neural networks. Neural computation , 1(2):270-280, 1989.

Wu, S. Yin, W. Qi, X. Wang, Z. Tang, and N. Duan. Visual chatgpt: Talking, drawing and editing with visual foundation models. arXiv preprint arXiv:2303.04671 , 2023.

Wu and K. He. Group normalization. In Proceedings of the European conference on computer vision (ECCV) , pages 3-19, 2018.

Wu, G. Tucker, and O. Nachum. Behavior regularized offline reinforcement learning. arXiv preprint arXiv:1911.11361 , 2019.

Zheng, A. Zhang, and A. Grover. Online decision transformer. In International Conference on Machine Learning , pages 27042-27059. PMLR, 2022.

Zheng, M. Henaff, B. Amos, and A. Grover. Semi-supervised offline reinforcement learning with action-free trajectories. In ICML , 2023.

Zhu, M. Dang, and A. Grover. Scaling pareto-efficient decision making via offline multi-objective rl. In International Conference on Learning Representations , 2023.

Appendices

A Model details of Decision Stacks

In Decision Stacks, we employ a modular approach by utilizing separate models for predicting observations, rewards, and actions. These models are trained independently using the technique of teacher forcing [Williams and Zipser, 1989]. This modular design allows us to explore the architectural flexibility of Decision Stacks and incorporate various inductive biases tailored to each modality. Below, we provide a description of the models used for each module in our approach:

A.1 Observations models

Diffusion-based model. We adopt the diffusion process as a conditional generative model to generate future observations based on current and past observations. Following Decision Diffuser [Ajay et al., 2022], we employ classifier-free guidance [Ho and Salimans, 2022] and low-temperature sampling to generate future observations conditioned on goals. With classifier-free guidance, we first sample x K ( τ ) from a Gaussian noise and refine it to x 0 ( τ ) with the following procedure in Eq 7 for denoising x k ( τ ) into x k -1 ( τ ) .

where G is the goal information. For the observation model, we diffuse over a sequence of consecutive observations.

We parameterize the diffusion model as U-Net [Ajay et al., 2022, Janner et al., 2022]. To generate future observations, we use an inpainting strategy where we condition the first observation in the diffusion sequence as the current observation o 0 .

Transformer-based model. The trajectory sequences can also be autoregressively predicted using a transformer architecture as previously demonstrated in Decision Transformer [Chen et al., 2021] and Trajectory Transformer [Janner et al., 2021]. In our approach, we employ a GPT [Radford et al., 2018] model as the underlying transformer architecture. This transformer-based model is trained by feeding it with a sequence of observations, alongside the task goal as the conditioning information. To ensure temporal consistency during both training and testing, we employ causal self-attention masks, enabling the model to predict the next observation solely based on the preceding observations and task information. To effectively represent the observation and task modalities, we treat them as separate entities and embed them into dedicated tokens. These tokens are then input to the transformer, which is further enriched with positional encoding to capture the temporal relationships among the observations. This combination of a GPT model, self-attention masks, and modality-specific token embeddings allows our approach to effectively model the sequential nature of the observations while incorporating task-related information. The transformer can autoregressively generate a sequence of observations as in Eq 8.

A.2 Reward models

Diffusion-based model. For our reward model, we employ a similar U-Net architecture as the observation diffusion model but diffuse over different sequences. Specifically, we diffuse over the combined sequence of observations and rewards as follows:

In this combined sequence, we concatenate the observation sequence o 0 , o 1 , ..., o T and the corresponding reward sequence r 0 , r 1 , ..., r T at each time step for training. To generate the reward sequence, we utilize an inpainting conditioning strategy, similar to the one used in the observation model. This strategy involves conditioning the diffusion process on the observation sequences o 0 , o 1 , ..., o T while generating the rewards. By incorporating this inpainting conditioning, the reward model can effectively utilize the available observation information to generate accurate reward predictions throughout the diffusion process.

Transformer-based model. We employ an Encoder-Decoder transformer architecture for reward sequence generation, given a sequence of observations and the task goal. The encoder module embeds the observations and task goal, incorporating time encoding for capturing temporal dependencies. The encoded inputs are then processed by a transformer layer. The decoder module generates the reward sequence by iteratively predicting the next reward based on the encoded inputs and previously generated rewards. The transformer architecture facilitates capturing long-range dependencies and effectively modeling the dynamics between observations and rewards.

MLP reward model. The multi-layer perceptron (MLP) reward model is a straightforward mapping from the current observation to the corresponding immediate reward. This model does not incorporate context or future synthesized information as it relies on a fixed input and output size. Consequently, the MLP architecture does not consider or incorporate any contextual information during its prediction process.

A.3 Action models

Diffusion-based model. Similar to the observation and reward diffusion processes, we perform diffusion over the trajectory defined in Equation 10 as well as goal conditioning. By employing diffusion in this manner, we generate the action sequence conditioned on the information contained within the previous observations, rewards, and anticipated future observations and rewards. This diffusion process enables us to effectively capture the dependencies and dynamics among these elements, resulting in the generation of contextually informed action sequences.

Transformer-based model. We employ the Encoder-Decoder transformer architecture, similar to that utilized in the reward transformer model, for our action model. The encoder module performs embedding of the observations, rewards, and task goals, incorporating time encodings to capture temporal dependencies. The encoded inputs are subsequently processed by a transformer layer. The decoder module is responsible for generating the action sequence by iteratively predicting the subsequent action based on the encoded inputs and previously generated actions. Consequently, our action model produces coherent and contextually informed action sequences.

MLP action model. The multi-layer perceptron (MLP) action model functions as a direct mapping from the current observation and subsequent observation to the corresponding immediate action, similar to the inverse dynamics model in Decision Diffuser [Ajay et al., 2022]. However, due to its fixed input and output size, this MLP architecture does not incorporate contextual or future synthesized information. Consequently, the model lacks the ability to consider or integrate contextual details. This limitation proves disadvantageous in the context of partially observable Markov decision processes (POMDPs), where the inclusion of contextual information is vital for inferring hidden states and making informed decisions.

B Hyperparameters and training details

The diffusion-based observation model follows the same training hyperparameters as in Decision Diffuser [Ajay et al., 2022], where we set the number of diffusion steps as 200 and the planning horizon as 100.

For other models, we use a batch size of 32, a learning rate of 3 e -4 , and training steps of 2 e 6 with Adam optimizer [Kingma and Ba, 2015].

The MLP action model and the MLP reward model is a two layered MLP with 512 hidden units and ReLU activations.

The diffusion models' noise model backbone is a U-Net with six repeated residual blocks. Each block consists of two temporal convolutions, each followed by group norm [Wu and He, 2018], and a final Mish nonlinearity [Misra, 2019].

For Maze2D experiments, different mazes require different average episode steps to reach to target, we use the planning horizon of 180 for umaze, 256 for medium-maze and 300 for large maze.

For Maze2D experiments, we use the warm-starting strategy where we perform a reduced number of forward diffusion steps using a previously generated plan as in Diffuser [Janner et al., 2022] to speed up the computation.

The training of all three models, including the observation, action, and reward models, is conducted using the teacher-forcing technique [Williams and Zipser, 1989].

Additional hyperparameters can be found in the configuration files within our codebase.

Upon reproducing the Decision Diffuser (DD) approach [Ajay et al., 2022] using their provided codebase, we observed that the agent performance can be sensitive to the test return of the locomotion tasks, conditional guidance parameter, and sampling noise. For Decision Stacks variants that use the DD observation model, we directly use the models tuned for DD's best performance. In future work, it can be helpful to use conservative regularizers [Nguyen et al., 2022] to further improve both DD and DS performance.

Each model was trained on a single NVIDIA A5000 GPU.

C Sensitivity analysis on the dimension occlusion for POMDPs.

D Ablation of reward modeling across MDP and POMDP.

We performed an ablation study on reward modeling, focusing on six locomotion environments. We trained two sets of transformer-based action models: one set modeling the sequence of observations and actions, and the other set incorporating the sequence of observations, rewards, and actions. The evaluation of these models was conducted on both Markov Decision Process (MDP) and Partially Observable Markov Decision Process (POMDP) environments, specifically three locomotion tasks. The results obtained from this study highlight the evident advantages of reward modeling, particularly in dense-reward locomotion tasks. The details and findings are presented in the Table 8, emphasizing the benefits gained from incorporating reward modeling.

E Modeling ordering ablation study

Table 9: Performance Comparison on Modeling Ordering. From the chain rule of probability, any autoregressive factorization can model the data distribution under idealized conditions. In practice, we choose the ordering of observations, rewards, and actions. Specifically, we ordered observations prior to rewards to be consistent with the functional definitions in MDPs, where the reward is typically a function of observations (and potentially other variables), but not vice versa. We tested this choice empirically as well in early experiments, and found our choice to significantly outperform the counterpart. The table provides a comparison of two different orderings: Reward-State-Action (R, S, A) and State-Reward-Action (S, R, A). The results indicate that the S, R, A ordering outperforms the R, S, A ordering in the halfcheetah-medium-replay-v2 environment. DS models observations prior to rewards to be consistent with the functional definitions in MDPs, where the reward is typically a function of observations (and potentially other variables), but not vice versa.

F Maze2D additional analysis

We investigate trajectory sampling-based methods, namely DD, Diffuser, and Decision Stacks, within this specific environment. All three models utilize a diffusion-based approach to generate future plans. In this task with goal-conditioning, the position of the maze's goal serves as the condition.

Regarding Diffuser, we strictly adhere to their codebase and observed that Diffuser impaints the last goal position into the diffusion sequence. Similarly, for DD, we follow their codebase and incorporate the goal as conditioning, embedding it and transmitting it to the diffusion backbone, which is an Unet. In the case of DD, we also apply inpainting conditioning by conditioning the last position in the diffusion sequence to serve as the goal position.

F.1 Detour problem of inpainting conditioning

However, we encountered an issue with this inpainting goal conditioning when performing open-loop generation. The problem arises because, during each replanning iteration of the diffusion model, it only conditions on the last time step to be the goal state. Consequently, when the agent is in close proximity to the goal point, the diffusion model plans a detour that initially takes the agent away from the goal and then redirects it back towards the goal. Since the environment's reward function is designed in such a way that the agent receives a reward when it is near the goal, and the episode does

not terminate upon reaching the goal point but rather when it reaches the maximum episode length, it is actually more advantageous for the agent to remain at the goal point.

After recognizing this issue, we devised a "progressive conditioning" (PC) strategy to address it. This approach involves gradually increasing the number of timesteps in the diffusion sequence that are conditioned to be the goal position as time progresses. By implementing this progressive conditioning method, we successfully resolve the detour problem and provide additional incentives for the agents to remain at the goal position. Without PC, the agent's behavior during open-loop evaluation exhibits a recurring pattern of moving toward the goal and deviating from it. This results in unfavorable trajectories where the agent fails to stay at the goal position.

The integration of progressive conditioning led to performance enhancement in both Diffuser and DD, although the overall results still indicate that Decision Stacks outperforms DD in both single goal (average performance: 109.8 vs. 131.5) and multi goal environments (average performance: 111.6 vs. 123.4). We compare the performance of Diffuser, DD, DS with and without progressive conditioning in the bar charts below in Figure 4 and Figure 5. The results indicate that DS with PC outperforms other baselines on most of the environment settings across single-goal and multi-goal settings.

Single Goal setting

199

179

159

139

119

99

79

59

39

19

1

Diffuser w/o PC Diffuser w/ PC

DD w/o PC

DDw/PC

DS w/o PC

DS w/ PC

umaze

medium large

MultiGoalsetting

159

139

119

99

79

59

39

19

1

Diffuser w/o PC Diffuser w/ PC

DD w/o PC

DD w/PC

DS w/o PC

DS w/ PC

umaze

medium large

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper focuses on solving offline RL with generative models. Concretely, it proposes a method to modularize the joint distribution of time-induced trajectories and use separate generative models to represent observation module, reward module, and action module. Evaluations are conducted on D4RL benchmark with MDP and POMDP environments. Extensive comparisons against prior works are included.\n\nSTRENGTHS:\n- The paper is fairly well presented and easy to follow. \n\n- The proposed method is extensively evaluated against multiple related approaches.\n\nWEAKNESSES:\n- Novelty of the proposed method is limited. The improved performance can also be attributed to other confounding factors, such as larger models due to reward and action modules being factored out.\n\n- It is counter-intuitive to ignore the canonical time-induced casual ordering in favor of different token types. More intuitions and theoretical analysis (if applicable) are encouraged to provide.\n\n- Lack of evidences to support the claims on modular expressivity (L145). Are there any experiments showing it can transfer to new environments?\n\n- The empirical results, especially those on D4RL locomotion tasks, are not significant enough to justify the extra compute introduced.\n\n- Just swapping each module with different model architectures/modeling strategies is not enough to gain insights of the proposed method. More ablations and analysis are necessary.\n\nQUESTIONS:\n- Why DS's performance on POMDP Hopper task with medium-expert data (Table 3) is even better than that on fully observable setting (Table 2)?\n\n- How does DS compare to more straightforward baseline, such as the one with separate heads for three modules but shares the same backbone?\n\nLIMITATIONS:\n- The design choice of ignoring temporal causality is not well motivated. Case studies or even proof-of-concept experiments would be helpful to justify it.\n\n- The performance improvement could be attributed to confounding factors such as larger model sizes. More experiments with controlled variables (e.g., same amount of model parameters) would help to support claims made in the paper.", "TITLE": "Official Review", "RECOMMENDATION": 3, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-06-28"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper highlights a drawback in prior frameworks, such as Decision Transformer and Diffuser, where the absence of modular hierarchies among different tokens results in limited expressivity and flexibility. To overcome this issue, the paper introduces Decision Stacks (DS), a modular algorithm designed for learning goal-conditioned policies using offline datasets. The proposed method parameterizes three generative model-based modules for future observation prediction, reward estimation, and action generation, respectively. Through various offline evaluations in both MDP and POMDP environments, this paper demonstrates that DS outperforms previous approaches by generating superior plans.\n\nSTRENGTHS:\nThis paper presents a modular probabilistic framework, utilizing deep generative models to establish token-level hierarchies in trajectory generation. The proposed algorithm, Decision Stacks (DS), incorporates independent generative models for simulating the temporal evolution of observations, rewards, and actions, allowing for parallel learning and enabling flexible generative decision making. The algorithm is both straightforward and effective, particularly in the POMDP setting, where DS surpasses other baseline methods by a significant margin. Additionally, the experiments are conducted meticulously, and the visualizations of example rollouts in the Maze2D-medium-v1 environment are clear.\n\nWEAKNESSES:\nOne major concern regarding the proposed algorithm DS is the high complexity of the model, as it requires training three generative models. However, in most experiments such as offline RL with an MDP setting, the performance improvement compared to other baselines that use only one generative model is not significant. Furthermore, the paper lacks an explanation of the specific dimensions that were excluded to construct the POMDP setting. Additionally, as depicted in Figure 1, the generation of observation, reward, and action sequences follows a sequential order, and there exist dependencies among the three generative models. However, the paper does not provide clear explanations on how training can be parallelized and how the choice of generative models affects performance.\n\nQUESTIONS:\n(1)How is the handcoded controller implemented in 4.1 and 4.2, respectively? What role does it play? In Table 1 and Figure 1, why do Decision Diffuser (DD) and DS exhibit similar performance when this controller is added? DD does not adopt a modular structure but utilizes the same diffusion-based observation model as DS. Does this indirectly suggest that the modular structure has minimal impact as long as a good generative model is chosen?\n(2)During the experiment, why was a diffusion model used for the observation model while transformer models were used for the action and reward models? Why do different choices of generative models lead to the results shown in Table 4? It is hoped that the author can provide corresponding reasonable explanations instead of simple combination attempts.\n\ntypoes, e.g., \"seggregate\" --> \"segregate\".\n\nLIMITATIONS:\nMethods such as Decision Transformer [1], which reduce reinforcement learning to a prediction task, have gained popularity partly due to their simplicity. In contrast, DS trades a high model complexity for relatively modest expressivity. It would be more convincing in terms of expressivity if experiments were conducted in stochastic environments[2] that yield better results. Furthermore, the paper lacks rational analysis in certain aspects, such as the selection of generative models to generate higher quality trajectories.\n[1] Chen L, Lu K, Rajeswaran A, et al. Decision transformer: Reinforcement learning via sequence modeling[J]. Advances in neural information processing systems, 2021, 34: 15084-15097.\n[2] Paster K, McIlraith S, Ba J. You Can't Count on Luck: Why Decision Transformers Fail in Stochastic Environments[J]. arXiv preprint arXiv:2205.15967, 2022.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-03"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper proposes to tackle the problem decision making using a stack of different generative models. The first generative model constructs a conditional distribution over observations. A subsequent generative model constructs a conditional distribution over rewards, with a final generative model constructs a distribution over actions. The authors illustrate the efficacy of this decomposed approach across a suite of different tasks.\n\nSTRENGTHS:\n- I enjoyed reading the paper -- it was quite clear and easy to read. The motivation of the paper to decompose the generative modeling objective into a set of component modules is sound. \n- The paper follows a formulation of offline reinforcement learning as probabilistic inference that is likely to relevant to a large audience at NeurIPS with the increasing popularity of generative models\n- The approach performs well across a set of different environments, outperforming existing baselines across 3 separate tasks.\n\nWEAKNESSES:\n- The results illustrated in the approach are a bit toy -- with the largest gains in the Maze2D environment. The Mujoco control environment has somewhat limited gains and the constructed POMDP environment seems a bit artificial. It may be that the Mujoco control environments have already saturated in performance and it may be interesting to try more complex planning tasks such as RLBench or other robotic manipulation environments.\n- While I understand the motivation to decompose trajectory synthesis into a set of modular components, I'm not sure I completely understand the particular decomposition of first observations, then rewards, and then actions. It seems like decomposing rewards first may be more natural then observations.\n- It might also be interesting to explore the extent to which decomposed submodules can enable compositional generalization or more efficient multitask learning by encoding structure in the greneration procedure. For instance -- perhaps the observation model can be used across a set of different tasks.\n\n\nThere are a couple of typos in the paper listed below:\n\n- L133 typo auotregressive -> autoregressive\n- In L303  Dai et al should be Du et al. The method does not use a state model as a pretrained text2image model but rather a learned text2video model.\n- The table before section 4.2 is misformatted and extends to the margin of the paper\n\nQUESTIONS:\nI had a couple questions about the underlying approach that might be interesting to discuss:\n\n- In equation 6, we sequentially sample from each factored distribution to sample from equation 3. In practice, would we instead want to sample marginalized distribution over actions across all possible observations and rewards given the goal?\n- It might be interesting to think a bit about how the individual decomposed modules can also provide feedback with respect to each other. For instance a planned set of states may induce a poor distribution over rewards. Having the conditional reward model then provide some feedback to regenerate the planned set of states may then be interesting. [1] may be an interesting read\n\n[1] Composing Ensembles of Pretrained Networks via Iterative Consensus\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-05"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper proposes to disentangle the different modalities (reward, observations, and actions) utilized in offline goal-conditioned reinforcement learning instead of the standard temporal single-module structure. The paper contributes three independent generative modules that have the benefit of parallelizable training. The paper includes empirical results showing the proposed algorithm's comparable or superior performance to baselines on several MDP and POMDP environments.\n\nSTRENGTHS:\n1. The core idea to utilize a different decomposition of tokens to enable independent training of generative models appears interesting and novel. I enjoyed the simplicity of the approach, accompanied by the good empirical results. \n\n2. Overall, I found the paper to be well-written and clear. I liked that the authors are clearly familiar with the relevant related literature and used diagrams throughout to help illustrate concepts. \n\n3. Given the good empirical results, the idea of semantic decomposition (noted in strength 1) could be generally impactful for communities interested in transformers + RL.\n\n4. I appreciate the extensive experiments performed to support the claims of flexibility and performance.\n\nWEAKNESSES:\nIn general, I liked this paper. A few things that I think could improve it:\n\n1. I would have liked to see the training cost comparison to the other algorithms. \n\n2. In general, I would also like to see a more extensive discussion of the experimental results, as noted in questions 4-6 in the Questions section below. \n\n3. I feel like some of the claims are a bit strong. For example, how does the work \"guarantee\" expressivity and flexibility? This language is often used when we have an assured property of the system or algorithm, but I don't see any rigorous theoretical evaluation of this claim. \n\nMore minor:\n\n1. Missing some references, including work in compositional offline RL (as mentioned in the introduction, first paragraph) [Mendez et al., 2022] and a reference to the POMDP formalism [Kaelbling, Littman, and Cassandra, 1998]. \n\n2. The focus on MDP and POMDP environments was somewhat unclear to me. I think including some motivation or intuition in the Introduction as to why this approach would work better than previous works in both types of environments would make this more clear.\n\n-----------------------------------------------------------------------------\n\n[Mendez et al., 2022] Mendez, Jorge A., Harm van Seijen, and Eric Eaton. \"Modular lifelong reinforcement learning via neural composition.\" arXiv preprint arXiv:2207.00429 (2022).\n\n[Kaelbling, Littman, and Cassandra, 1998] Kaelbling, Leslie Pack, Michael L. Littman, and Anthony R. Cassandra. \"Planning and acting in partially observable stochastic domains.\" Artificial intelligence 101.1-2 (1998): 99-134.\n\nQUESTIONS:\n1. Recent work [Carroll, et al. 2022] incorporates dynamics learning in some training setups and demonstrates improved performance over vanilla Decision Transformer. Is this not included in the experimental comparisons due to the different training scheme & architecture choice? \n\n2. Why do only some of the algorithms have a reported error (e.g., in Tables 1&2)? Also, is the value standard error, standard deviation?\n\n3. Could the authors please clarify what is meant by the expert-normalized scores in Section 4.2? \n\n4. What is the training cost of the proposed algorithm (in terms of both samples and wall clock time)? How does it compare to the baselines?\n\n5. Why is the architectural flexibility investigated using Hopper-medium-v2 despite it not being used in the other experiments (e.g., those presented in Table 3)? \n\n6. Why does DS work better on medium but not medium expert or medium replay HalfCheetah (Table 3)? \n\n7. How were the two dimensions of the state chosen to exclude (Section 4.3)? What is the observation representation in these experiments?\n\n---------------------------------------------------------------------------------\n[Carroll et al., 2022] Carroll, Micah, et al. \"Uni [mask]: Unified inference in sequential decision problems.\" Advances in neural information processing systems 35 (2022): 35365-35378.\n\nLIMITATIONS:\nThe authors have a small section in the Conclusion focusing on the limitations. I felt like this was sufficient. However, I can always appreciate a lengthier discussion section, especially when one's proposed algorithm performs similarly to other algorithms in some settings.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-09"}

comments: SUMMARY:
The paper proposes using different generative models (transformer or diffusion) rather than the same model (like in trajectory transformer/decision diffuser) for observation prediction, reward estimation, and action prediction in “model-based” offline RL. They demonstrate that this flexibility improves performance in offline RL in a POMDP setting where two dimensions of the state vector are removed for each environment.

STRENGTHS:
1. The paper is well written and easy to read.
2. The consistent increase on POMDPs is encouraging. And the ablation with different modules is very interesting.

WEAKNESSES:
1. Similar performance to DD on D4RL Gym tasks.

2. Major similarities to previous work (trajectory transformer, decision diffuser) and the authors acknowledge this and highlight the reward modeling as novel. While it is interesting, a (temporal difference) variant already exists in the trajectory transformer paper [Janner et al 2021].  Namely, the trajectory transformer with Q function (a form of reward modeling) guided planning which can be found in Section 4.2 of [Janner et al 2021], under “Combining with Q functions”. 

3. Does it matter which two observation dimensions are removed? An explanation of which dimensions were removed and why is missing and a sensitivity analysis to the dimensions removed are missing as well.

4. Is there any intuition for why the conditioning in Equations (4), (5), (6) were chosen? 

5. Is there heuristic that could be used to choose the different modules with any online interaction? Such a heuristic would be necessary in offline RL.

QUESTIONS:
Please see weaknesses.

LIMITATIONS:
Yes.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-28