Decision Stacks: Flexible Reinforcement Learning via Modular Generative Models
Authors: Siyan Zhao
Aditya Grover
Venue: NeurIPS 2023 poster (2023)
Keywords: reinforcement learning
generative models
offline RL
sequential decision making
TL;DR: Decision Stacks is a generative framework for goal-conditioned RL that models trajectories as composable modular components for observations, rewards, and actions leading to superior performance for planning and RL in MDP and POMDPs.

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
Siyan Zhao
Aditya Grover
Department of Computer Science University of California Los Angeles siyanzecs. ucla. edu
Department of Computer Science University of California Los Angeles adityagecs. ucla. edu
Abstract
Reinforcement learning presents an attractive paradigm 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, generative framework that decomposes goal-conditioned policy agents into 3 generative mod- ules. 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, optimiza- tion 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
Introduction
Modularity is critical design principle for both software systems and artificial intelligence (AI), It allows for the creation flexible and maintainable systems by breaking them down into smaller, inde- pendent 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 al. [203] defines family of 22+ vision and language foundation models, such as ChatGPT [OpenAl] 2022 (language generation). SAM [Kir illov al. 2023 (image segmentation). and StableDiiffusion [Rombach et al. [002 (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 spatiotemporal abstractions that simplify complex goals via hierarchical RL, e.g., [MGGovern and Barto 2001 Andreas et al. 2017 Simpkins and Isbell] 2019 Ahn et al. 2022 Kulkarni et al. 2016 Mendez et al. [021]. Distinct but complementary
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 multitude of different tokens of information: goals, observations, rewards, actions. As shown in many recent works (Chen) al. [222] Janner al. 2021 2022 Ajay al., 2022 Zheng et 2022 Reed et al.) 2022), can effectively reduce RL to probabilistic inference
Levine 2018) via learning deep generative models
'The project website and code can be found here: http:///iiaazzhaaaoggiibbiiiiiii
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

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, family of generative algorithms for goal-conditioned RL featuring 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 [ for an illustration. While our factorization breaks the canonical time-induced causal ordering 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 (c.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 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 tradeoff' 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 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, A, G, P, (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 c.g., a language instruction or a visual destination designed cumulative return for single-task policies. The transition probability function, P S x A x S > [0, 1]. describes the transition dynamics. Meanwhile, receives after performing action a in state 8. The observation emission model & = P(o]s), determines the probability of observing in state s. The 7. Po (so). 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 coal-conditioned RL objective is to find the optimal policy x² that maximizes the expected cumulative discounted reward over the episode
state for multi-task policies,
R: 9 x S x A > I defines the rewards that the decision-mall
horizon: 7M := EGGgg (ee
Σi=o ''r (&t: 0e Gi)
2

---

+
0.01
0.28
f
..41
o
{
G, Oo
ro
a
01
r1
Q1
02
T2
a2
OT
TT
aT
}
Goal G
Initial Observation og
P(onT 00, G)
P(ror our G} Reward Estimation
P(aor 00:7, oo:,, G)
Observation Prediction
Action Generation
0.01
0.28
0.41
0.98
Figure I: Illustration for the Decision Stacks framework for learning reinforcement learning agents using probabilistic inference. In contrast to a time-induced ordering, we propose a modular design that segregates the modeling of observation, rewards, and action sequences. Each module can be flexibly parameterized via any generative model and the modules are chained via an autoregressive dependency graph to provide high overall expressivity.
2.2 Offline reinforcement learning
Offline reinforcement learning (RI. a paradigm for policy optimization where the agent is only given access to 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 challenging or infeasible such as healthcare [Murphy et al.) [2001] and autonomous driving. A major obstacle in offline RL is dealing with distributional shifts. we naively use Bellman backup for learning the Q-function of given policy, the update which relies on the actions sampled from policy л can learn inaccurately high values for outto--distiibution actions, leading to instability in bootstrapping
process and causing value overestimation [Kumar al. 2020
2.3 Generative models: Autoregressive transformers and Diffusion models
In this work, we are interested in learning the distribution Pdata (x/c) using a dataset D consisting of trajectory samples and conditioning We consider two conditional generative models for parameterizing our agent policies to learn the distribution:
Transformer is powerful neural net architecture for modeling sequences Vaswani et al. 2017 It consists of multiple identical blocks of multi-head self attenttnnio modules and position- wise fully- connected networks. The vanilla transformer can be modified with causal self-attention mask to parameterize an autoregressive generative model as GPT [Radford al. (0188). Autoregressive generative models, such as the transformers, factorize the joint distribution (( x.)) product of conditionals, which can represented as: p(x) = II p(x,)xx.)) This equation shows that the probability of each variable r; depends on the previous variables 21. Ti-1. One advantage of this factorization 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 v; sampled based on its preceding variables.
Diffusion Models (Sohl-Dickstein et al. 2015 Ho et al. 2020 are latent variable models that consist of predefined forward noising process g (P 2t == (a++ii Vorgar, (1 ou (v)) that gradually corrupts the data distribution q(æn)) into (0, 7) in K steps, and a learnable reverse denoising process po -1 (ma) = (\a (1 k). x.). For sampling, we first generate latent sample from the Gaussian prior NO0 I) and gradually denoise it using the learned model pe (#x-1 1 for

---

Figure 2: Graphical model for the data genera- tion process in a POMDP. Here, we show the case where the behavioral policy based on hidden state information. The dashed cir- eles implies that this state information is not stored in the offline dataset. G represents the task condi- tioning, e.g., a target return (for single-task agents) or navigation goal (for multi-task agents).
potentially act
K steps to obtain the data sample 2o. Diffusion models can be extended for conditional generation using classifier-free guidance [Ho and Salimans] 2022] where conditional eg (#rg, c, k) and an unconditional eo (ank. k) noise model is trained. Conditional data sampled with the perturbed noise <g (#t, k) + (ce (ark. e, k)) ca (#g, k)), where w is the guidance strength and is the conditioning information.
Flexible and Modular RL via Decision Stacks
In Figure we consider 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( 0z), goal-dependent reward function P(re(Se, dg: G), and the observation emission probability Ps).(s:). Unlike the learning agent which only has access to observations, the behavioral policy used for generating offline trajectory dataset might also have access to the hidden state information. Such situation is common in real-world applications, e.g.. a human demonstrater may have access to more information about internal state than what is recorded in video demonstration datasets available for offline learning. Finally, note that Markov Decision Process (MDP) can be viewed as a special case POMDP, where the observation at each timestep, or, matches the underlying state, St. To avoid ambiguity, overload the use of to denote both the state and observation in an MDP at time
In the context of goal-conditioned decision-makinn a finite-horizon trajectory in the offline dataset D is composed of goal G and sequence of observations, actions, and reward tokens.
(G, ou, ag To: ..... O₁, ag ........ OT, ar, ar,
(1)
Our primary objective lies in learning a goal-conditioned distribution for (ao.7 01.7, Torrlon, G) conditioned on an arbitrary goal G and initial observation og O. Leveraging the chain rule of probability, we can factorize this joint distribution into product of conditional probabilities. For example, Janner al. [202T] use time-induced autoregressive factorization:
T
Paata (ao:T, 01.T, To:r 0o. G)
II Po(odTT Po T<D II Polaton,a (2)
where Tee denotes all the tokens in the trajectory before time t. Each conditional factor is parameter- ized via an autoregressive transformer with shared parameters 0. the parameterization sufficiently expressive, any choice of ordering for the variables suffices. However, in practice, we limited by the size our offline dataset and the choice of factorization can play a critical role.
In Decision Stacks, we propose to use modular factorization given
Pana (ao:T, 01:T, To:T Oo, G) xe Po. (01.T oo. G) Pa. (ro.T lo.r, G) Pe, (ao.T lo.T, To:7, G)
observation module
reward modulc
action module
Each of the modules (observations, rewards, or actions) focuses on predicting a distinct component of the POMDP and has its own set of parameters (01, 02, 0z). Our motivation stems from the fact that real-world domains, each component is sufficiently distinct from the others its semantics and representation. Such variances span across multitude of factors including dimensionalities, domain types (discrete or continuous), modalities (e.g., visual observations, numeric rewards), and information density (c.g., rewards can be sparse, state sequences show relatively high continuity).

---

Modular Expressivity. In Eq. 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 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 goal. For training a multi-task agent, the goal can be specified flexibly as spatial coordinates, 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 cenutt-conditoneed 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 [989]. At test-time, our framework naturally induces a planner, as in order to predict an action at time f, 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 likely to be more accurate as updates the inputs to the modules based on the environment outputs at each time-step. Using closed-loop plan, we can sample the action at time I as follows:
ōU+1:T ~ Po; (01+1:T 00:1, G)
(4)
f++1T ~ Poz (Tt+1.T To:t; Opst, 0c+1:T G)
(5)
à 1g~Pa Po. (ax 00 0 To:t, T',,,, ft+1:T G)
(6)
The hat symbol (*) indicates predicted observations, rewards, and actions, while its absence de- notes observations, rewards, and actions recorded from the environment and the agent in the pre- vious past timesteps. For closed-loop planning, Eqs. observation, reward and action distributions on the past trajectory tokens. For 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 observa- tion module is an autoregressive transfer, then we can obtain the sampling distribution in Eq. as:
require us to condition the joint
Po, (0t+1:7 0o:t; G) II Po, (o: des, 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 more detailed discussion of broad prior works in 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 al. 2021] and tiffusion-basse extensions (Wang et al., 2022). Second, there exist mwwem-ddddeeeeooooo approaches but make different design choices; Trajectory Transformer [Janner et al.) 202T] uses a lime-induced causal factorization parameterized by a single autoregressive transtormer, 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 based inverse dynamics model to extract actions. Unlike these works, we propose a modular structure that is maximally expressive as 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.
5

---

Table 1: Performance on Maze2D tasks. DS significantly outperforms other baselines without the need for handcoded controller. Note that DD and DS share the same diffusion-basee observation model architecture and hence with a handcoded controller, their performance is the same. We average the results over 15 random seeds and emphasize in bold scores within 5 percent of the maximum per task (2 0.95 max).
Diffuser with Handcoded Controller
DS/DD with Handzoded Controller
Tiask
Enviromment
MPPI
CQL
IQL
Diffuser
DD
D8
33.2
5.7
47.4
86.9
113.8
111.3
113,9
119.5 #26
Single Goal
medium
10.2
34.9
108.5 #174
103.7 #212
111.7 #2.4
121.5
112.9 alLa
large
5.1
12.5
58.6
45.4 =14.5
111.8 243.4
171.6 #13.
123.0 =6.4
132.8 $21.
Average
16.2
7.7
47.0
80.2
109.8
131.5
119.5
121.7
umaze
41.2
24.8
114.4
a163
105.6
129.3
=29.011
136.1
Mali Goals
modium
15.4
12.1
$1.6 =14.5
126.4 #14.5
122.3 #5.7
127.2 =5.4
124.6 =11.3
large
8.0
13.9
41,0=20.1
116.0 #:3.1
126.7 #21
132.1 20.8
134.8 =12.3
Average
70.0
123.4
129,4
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 offiin RL methods in MDP environments?
$4.3
How does Decision Stacks compare with other offline RL methods in POMDP environments?
8..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?
experiment with D4RL environments and parameterize Decision Stacks with a iiffusion-basee observation model, an autoregressive transformer- based reward model, and an autoregressive transformer-bast based action model. Finally, in $4.4. we will ablate the full spectrum of architecture design choices for each module.
For
and
4.1
Long-Horizon Coal-Conditioned Environments
We first test for the planning capabilities of Decision Stacks the Maze2D task from the D4RL Fu et al. 2020 benchmark. This is a challenging environment requiring an agent to generate a plan from start location to goal location. The demonstrations contain sparse reward signal of +1 only when the agent reaches close to the goal. Following Janner et al.) [2022] we consider 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 al.
2013). 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 al., 2022 and
Decision Diffuser (DD) [Ajay et al. [222].
We show our results in Table for different goal types and maze grids. While WhileJanner al. 2022 previously demonstrated remarkable ability in generating long-horizon plans using Diffuser, their trajectory plans were executed by 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 Figure We hypothesize this could stem from the lack of modularity in Diffuser affecting the generation lidelity, or the lack of expressivity in using an MLP-based inverse dynamics model in DD which limits the context length required for long- hong-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 high- dimensional 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 al., 202T [ and
CQL Kummar et al. 2020), model-based algorithm MOReL [Kidambi et al. 2020 transformer-base based

---

Goal
Diffuser
DD
DS
152
Diffuser with Handcoded Controller
DS/DD with Handcoded Controller
Example rollouts on the Maze2D medeium-e environment. The goal is located at the bottom right corner of the maze. The trajectory waypoints are color-coded, transitioning from blue to red as time advances. The bottom two rows demonstrates that Diffuser, DD, and DS are all capable of generating good plans that can be executed well with handcoded controller. However, the respective action models result in differing exccutions. Compared to DD and Diffuser, DS generates smoother trajectories that are more closely aligned with the future waypoints planned by the observation model.
Figure 3:
generative models such as Decision Transformer (DT) [Chen et al. 2021 and Trajectory Transformer (TT) [Janner et al. 2021]. and diffusion- basdd generative models Diffuser (Janner al. 2022 and Decision Diffuser (DD) [Ajay et al. [022]. 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, reproduced numbers are slightly lower. We provide more details the Appendix. For a fair comparison, we used the same of hyperparameters that give the best performance for the DD bascline.
We show results averaged over planning seeds and normalize the scores such that value of 100 represents 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 make appropriate conditional independence assumptions using the MDP framework, the expressivity 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 lack of relevant sensors. A sensitivity analysis on dimension occlusions further strengthens our results, as shown in Appendix Table DS continues to outperform other baselines for each environment from the D4RL locomotion datasets. We report our results in Table 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

---

Table 2: Offline Reinforcement Learning Performance in MDP. Our results averaged 15 random seeds. Following Kostrikov et al. 2021]. we bold all scores within 5 percent of the maximum per task (2 0.95 max).
DD (reproduced)
Enviromment
BC
IQL
MOReL
DD
Diffuser
(ours)
Medium. Expert
HaliCheetah
$5.2
66.7
91.6
86.8
95.0
53.3
90.6
91.5 =2.5
79.5
95.7 :0.5
Medium- Expert
Hopper
Mcdiun-Export
Walker2d
107.5
109.6
108.8
108.1
101.9
95.6
108.8
105.2-2.3
108.4
108.0 =0.1
Medium
HalfCheetah
Medium
Hopper
52.9
66.3
58.5
67.6
61.1
95.4
79.3
81.2 =7:2
58.5
76.6 #442
Medium
Walker2d
74.0
79.9
79.7
83.6
Medium Replay
HalfCheetah
36.6
44.2
45.5
166.6
41.9
40.2
39.3
39.4
42.2
41.1 =0.11
Medium- Replay
Hopper
82.7
95.3-3.7
Medium Replay
Walker2d
26.0
73.9
77.2
66.6
82.6
49.8
75
72.3=3.1
61.2
80.7
Average
74.7
78.9
82.2
80.3
75.3
81.1
most environments and attains the highest average score of 74.3, which reflects 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.
Table 3: Offline Reinforcement Learning Performance in POMDP. Our results are averaged over 15 random seeds. Following Kostrikov et al. [2021]. we bold all scores within 5 percent of the maximum per task 0.95 max).
Dataset
Environment
BC
DT
TT
DD
Diffuser
DS (ours)
Medium-Faee Eepiuum-Expert
HalfCheetah
42.1
80.8
94.9
19.07
82.2
92.7
Medium-Expent
Hopper
51.1
105.2
61.6
32.7
70.7
110,9 10.4
Medium-Expeet
Walker2d
51.3
106.0
51.7
74.8
82.4
94.1 +8.5
Medium
HalfChoctah
43.3
42.7
46.7
40.3
45.4
47.1 10..3
Medium
Hopper
36.4
63.1
55.7
38.1
62.2
57.7 13.9
Medium
Walker2d
39.4
64..2
28.5
53.2
55.7
+1.2 74.3
Medium
Medium-Replay
HalfChectah
2.1
35.5
43.8
39.8
39.3
40.3 11.2
Medium-Replay
Hopper
24.3
78.3
84.4
22.1
80.9
86.9 12.66
Kedium-RRplay
Walker2d
23.8
45.3
10.2
58.4
58.7
66.8 1118
Average
34.9
47.9
53.0
42.0
64.2
74.3
4.4 Architectural Flexibility and Compositional Generalization
Table 4: Performance on Hopper-medium-yy POMDP using various reward and action models, with diffusion- based transformer-bant observation model. In each choice of observation model, the algorithm with the highest performance is highlighted.
Action models
Reward models
Transformer
Diffusion
MLP
Transformer
Diffusion
MLP
Transformer
57.7 +3.9
58.2 #4.3
45.6 #4.1
53.0 #3.7
54.3 #3.3
36.7 #4.2
Diffusion
51.7 11.7
56.9 42.2
36.3 13.1
58.0 14.4
46.9 13.7
34,9 43.5
MLP
56.0 +3.5
52.6 #2.5
33.3 #3.0
55.0 #3.9
52.1 +2.7
42.5 #4.1
Diffusion-based observation model
Transformer-based observation model
Decision Stacks distinctly separates the prediction 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 MILPs that are restricted in their window and only look at the immediate state information to make decision. The

---

Table 5: Ablation results comparing the performance of action models with and without reward information as input, across different architectures, in the dense reward POMDP task of Hopper- medium-v2. The results suggests clear advantage when incorporating reward modeling.
Action models
Action models
without reward modelling
with reward modelling
Observation model
Transformer
Diffusion
MLP
Transformer
Diftusion
MLP
Diffusion-basee
43.6 +1.3
43.7 #3.4
38.1 #2.1
57.7 #3.9
58.2 #4.3
45.6
#4.1
Transformer-baie
45.1 45.2
39.4 #3.2 39.6 4.3.7
58.0 #4.4
54.3 #3.3
42.5 #4.1
display 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 (state) + (reward) + (action) models. In Table 4
results shown in Table 4
we find that the performance of pure transformer- diffusion-baase 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 reward model and a diffusion- based action model. MLPs generally are outperformed by generative architectures, especially when used for modeling actions.
uses diffusion-baan
obsevation model, a transformer-
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 modules new environments. environment, altering action and reward spaces while keeping
shown in table we created modified
versions of the Maze2d-Umaze 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 -ineedd-mazeev2 environment, we analyze both dense and sparse rewards. For the action space, we consider three variations of the force on 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.
Action Space
Unsculed Force
Reflected Force
Rotated Force
dense
93.2 + 10.7
90.3 9.1
102.9 + 8.1
Reward Space
sparse
111.3 12.2
102.6 : 3.9
8.0
Furthermore, compare the best reward modeling architectures with alternatives that do not consider rewards. This standard practice for Diffuser [Janner et al.) 2022) and Decision Diffuser (DD) (Ajay et al. 2022]. For example, DD predict the action time only based the current observation and next observation, P(a lo. 01+1), parameterized via an MLP. As delineated in Table 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 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 value function that strictly
serves as lower bound for the true value function in CQL [Kumar al.) 2020). techniques focus
on uncertainty estimation such as Kumar et al. [2019]. and policy regularization methods [Wu
al. 2019 Ghasemipour et al. Kumar et al., 2019 Fujimoto and Gu 2021 Fujimoto et
2019 Model-based methods like MORel Kidambi et al., 2020] and ReBel. [ee et al. [021] add
pessimism into the dynamics models. In the context of partially observed settings, Rafailov et al. [202]] 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 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 RI. Over the past few years, the RI. community has seen a growing interest in employing generative models for context-conditionell sequence generation by framing the decision-mant making problem as a generative sequence prediction problem. Here, we expand on our discussion from $ with additional context and references. Decision Transformer [Chen et al. 2021 and Trajectory Transformer [Janner et al.) [021] concurrently proposed the use of autoregressive wansformer-baseer 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 sequence- level exploration strategy. GATO [Reed et al., 2022 Lee ct al. 2022 and PEDA [Zhu el. 2023) scale these models to multi-task and multi-objective settings. MaskDP Liu et al. [2222] shows that other super-sd objectives such as masking can also enable efficient offline RI,, especially in goal-conditioned settings. These ad advancements have also been applied to other paradigms in sequential decision making such black-box optimization and experimental design [guuee and Grover 2022 Krishnamoorthy et al.) 2223a 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) al. 202T Ajay et al. 2022 Wang al. 2022 Chen al. 2022). Dai 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 |
levell
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 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 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 Meta Research Award.
10

---

References
M. Ahn, A. Brohan, N. Brown, Y. Chehotar, O. Cortes, B. David, C. Finn, C. Fu, K. Gopalakrishnan, K. Hausman, A. Herzog, D. Ho, 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 can and not i say: Grounding language in robotic affordances. In arXiv preprint arXiv:2204.01691, 2022.
A. 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.1565,, 2022.
J. Andreas, D. Klein, and S. Levine, Modular multitask reinforcement learning with policy sketches. In International Conference on Machine Learning, pages 166-175. PMLR, 2017.
H. Chen, C. Lu, C. Ying, H. Su, and J. Zhu. Offline reinforcement learning via high-fidelity generative behavior modeling, arXiv preprint arXiv:2209. 2022.
L. Chen, 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.
Y. Dai, M. Yang, B. Dai, H. Dai, O. Nachum, J. Tenenbaum, D. Schuurmans, and P. Abbecl. Learning universal policies via text-guided video generation. arXiv preprint arXiv:2302.00111, 2023.
S. Emmons, B. Eysenbach, Kostrikov, and S. Levine. Rvs: What is essential for offline via supervised learning? arXiv preprint arXiv:2112. 2021.
J. Fu, A. Kumar, O. Nachum, G. Tucker, and S. Levine. D4rl: Datasets for deep data-driven reinforcement learning. arXiv preprint arXiv:2004.07219, 2020.
S. Fujimoto and S. S. Gu. A minimalist approach to offline reinforcement learning. Advances in neural information processing systems, 1222032-20144- 2021.
S. Fujimoto, D. Meger, and D. Precup. Off-policy deep reinforcement learning without exploration. In International conference machine learning, pages 2052-200200 25622 PMLR, 2019.
S. Ghasemipour, D. Schuurmans, and 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.
J. Ho and T. Salimans. Classifier-free diffusion guidance. arXiv preprint arXiv:2207.12595 2022.
J. Ho, A. Jain, and P. Abbeel. Denoising diffusion probabilistic models. Advances in Neural Information Processing Systems, 33:6840-6851) 2020.
M. 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.
M. Janner, Y. Du, J. B. Tenenbaum, and S. Levine. Planning with diffusion for flexible behavior synthesis. arXiv preprini arXiv:2205.09991, 2022.
L. 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.
R. Kidambi, A. Rajeswaran, P. Netrapalli, and T. Joachims. Morel: Model-based offline reinforcement learning. Advances in neural information processing systems, 33:21810-21823, 2020.
D. Kingma and J. Ba. Adam: A method for stochastic optimization. In International Conference on Learning Representations (ICLR), 2015.
A. Kirillov, E. Mintun, N. Ravi, H. Mao, C. Rolland, L. Gustafson, T. Xiao, Whitehead, A. C. Berg, W.-Y Lo, P. Dollár, and R. Girshick. Segment anything. arXiv:2004.02643, 2023.
11

---

I. Kostrikov, A. Nair, and S. Levine. Offline reinforcement learning with implicit q-learning, arXiv preprint 0rXiv:2110.06166 2021.
S. Krishnamoorthy, S. M. Mashkaria, and A. Grover. Generative pretraining for black-box optimiza- tion. In International Conference on Machine Learning (ICML), 2023a.
S. Krishnamoorthy, S. M. Mashkaria, and A. Grover. Diffusion models for black-box optimization. In International Conference on Machine Learning (ICML), 2023b.
T. 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.
A. 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.
A. Kumar, A. Zhou, G. Tucker, and S. Levine. Conservative q-learning for offline reinforcement learning. Advances in Neural Information Processing Systems, 33:1 13:1179-1191, 2020.
B.-J. Lee, J. Lee, and K.-E. Kim. Representation balancing offline model-based reinforcement learning. 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-gant gamt decision transformers. Advances in Neural Information 2022.
Processing Systems, 35:27921-27985
S. Levine. Reinforcement learning and control as probabilistic inference: Tutorial and review. arXiv preprint 1805 00909, 2018.
F. Liu, H. Liu, A. Grover, and P. Abbeel. Masked autoencoding for scalable and generalizable decision making. In Advances Neural Information Processing Systems (NeurIPS). 2022.
A. McGovern and A. G. Barto. Automatic discovery of subgoals in reinforcement learning using diverse density. In International Conference on Machine Learning (ICML), 2001.
J. A. Mendez, H. van Seijen, and E. EATON. Modular lifelong reinforcement learning via neural composition. In International Conference on Learning Representations, 2021.
D. Misra. Mish: A self regularized non-monotonic activation function. arXiv preprint 17X1v:1908.08681 2019.
S. 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, 16(455) 2001.
T. Nguyen and A. Grover. Transformer neural processes: Uncertainty- meta learning via sequence modeling. International Conference on Machine Learning (ICML), 2022.
T. Nguyen, Q. Zheng, and A. Grover. Reliable conditioning of behavioral cloning for offline reinforcement learning. arXiv preprint arXiv:2210005158, 2022.
OpenAI. Openai: Introducing chatgpt, 2022. URL htttps://openai.comm www.//og///aatg//////o/////////
M. L. Puterman. Markov decision processes: discrete stochastic dynamic programming. John Wiley
& Sons, 2014.
A. Radford, K. Narasimhan, T. Salimans, Sutskever, et al. Improving language understanding by generative pre-training. In preprint, OpenAl, 2018.
R. 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.
S. Reed, K. Zolna, E. Parisotto, S. G. Colmenarejo, A. Novikov, G. Barth-Maron, M. Gimenez, Y. Sulsky, Kay, J. T. Springenherg, Eccles, J. Bruce, Razavi, Edwards, Heess, Chen, R. Hadsell, O. Vinyals, M. Bordbar, and N. de Freitas. A generalist agent, 2022.
12

---

R. 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.
C. Simpkins and C. Isbell. Composable modular reinforcement learning, In Proceedings of the AAAI conference on artificial intelligence, volume 33, pages 9975-4982, 2019.
Sohl-Dickstein, E. Weiss, N. Maheswaranathan, and S. Ganguli. Deep unsupervised learning using nonequilibrium thermodynamics. In International Conference Machine Learning, pages 2256-2265. PMLR, 2015.
A. Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A. N. Gomez, L. Kaiser, and I. Polosukhin. Attention all you need. Advances neural information processing systems, 30, 2017.
Z. Wang, J. J. Hunt, and M. Zhou. Diffusion policies as an expressive policy class for offline reinforcement learning. arXiv preprint arXiv:2208.061 2022.
G. Williams, A. Aldrich, and E. Theodorou. Model predictive path integral control using covariance variable importance sampling. arXiv preprint 1509. 01 149, 2015.
R. J. Williams and D. Zipser. learning algorithm for continually running fully recurrent neural networks. Neural computation, 1(2):270-280, 1989.
C. Wu, Yin, W. Qi, Wang, Z. Tang, and N. Duan. Visual chatgpt: Talking, drawing and editing with visual foundation models. arXiv preprint arXiv:2303.0400 2023.
Y. Wu and K. He. Group normalization. In Proceedings of the European conference computer vision (ECCV), pages 3-19, 2018.
Y. Wu, G. Tucker, and O. Nachum. Behavior regularized offline reinforcement learning. arXiv preprint 1rXiv:1911.1133 2019.
Q. Zheng, A. Zhang, and A. Grover. Online decision transformer. In International Conference Machine Learning, pages 27042-27059. PMLR, 2022.
Q Zheng, M. Henaff, B. Amos, and A. Grover. Semi-supervised offline reinforcement learning with action-free trajectories. In ICML, 2023.
B. Zhu, M. Dang, and Grover. Scaling pareto-efficient decision making via offline multi-ohjective
rl. In International Conference on Learning Representations, 2023.
13

---

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 conditional generative model to generate future observations based on current and past observations. Following Decision Diffuser [Ajay et al., 2022], we employ classifier-fres guidance (Ho and Salimans, 2022] and onwtemperattue sampling to generate future observations conditioned on goals. With classifier-fre guidance, we first sample xx(r) from Gaussian noise and refine it to #0(
with the following procedure in Eq for denoising (Tx(T) into (x)1
é == En (T) k) + w (eo (#x G, (.x(T), k,))
where G is the goal information. For the observation model, we diffuse over sequence of consecutive observations.
Tabs 00
01
OT
We parameterize the diffusion model 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 oo-
Transformer-based model. The trajectory sequences can also be autoregressively predicted using transformer architecture previously demonstrated in Decision Transformer [Chen et al. 2021] and Trajectory Transformer [Janner et al.) 202T]. In our approach, we employ GPT [Radford al. [0188] model as the underlying transformer architecture. This transformer-bas model is trained by feeding it with a sequence of observations, alongside the task goal 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 GPT model, self-attention masks, and modality- token embeddings allows our approach to effectively model the sequential nature the observations while incorporating task-related information. The transformer can autoregressively generate sequence of observations in Eq g
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:
Ou
01
or
Trew
Fo
F1
Tr
In this combined sequence, we concatenate the observation sequence Ou, 01,..., Oy and the corresponding reward sequence Fo. ri,..., PT 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 og, 01, or 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.
14

---

OT Transformer-based model. We employ Siccoder-eccodee transformer architecture for reward sequence generation, given 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 straightforward mapping from the current observation to the corresponding immediate reward. This model does not incorporate context future synthesized information relies fixed input and output size. Consequently, the MIP 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 T0 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 the generation of contextually informed action sequences.
00
01
Tact
ro
T1
FT
(10)
@1 OT
Transformer-basee model. We employ the Encoder-Decoder transformer architecture, similar to that utilized 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 encoded inputs and previously generated actions. Consequently, action model produces coherent and contextually informed action sequences.
MLP action model. The multi-layer perceptron (MLP) action model functions 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 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 in Decision Diffuser [Ajay et al. [022]. where we set the number of diffusion steps as 200 and the planning horizon 1O0.
For other models, we use batch size of 32, a learning rate of 3e- and training steps of
6 with Adam optimizer Kiinmma and Ba 20155.
The MLP action model and the MLP reward model is 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 Ww and
He 2018), and final Mish nonlinearity [Misra] 2019
15

---

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-maz and 300 for
large maze.
For Maze2I) experiments, we use the warm-starting strategy where we perform reduced number of forward diffusion steps using previously generated plan as in Diffuser Danner
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-foreing technique (Williams and Zipser, 1989].
Additional hyperparameters can be found in the configuration files within codebase.
Upon reproducing the Decision Diffuser (DD) approach [Ajay et al., [022] using their provided codebase, observed that the agent performance be sensitive the test return of the locomotion tasks, conditional guidance parameter, and sampling noise. For Decision Stacks variants that 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 further improve both DD and DS performance. Each model was trained on single NVIDIA A5000 GPU.
C Sensitivity analysis on the dimension occlusion for POMDPs.
Table Sensitivity Analysis on Occluded Dimensions for POMDPs. In the Hopper environment, the full state contains 5 dimensions of position data and 6 dimensions of velocity data of different joints. The table below illustrates experiments with various dimensions and semantics for occlusion. DS consistently exhibits superior or second-best performance in comparison to other baselines on the pepper--edium-eexe
dataset. The standard error is reported, and the results are averaged over
15 random seeds.
Occluded Dimension
BC
DT
TT
DD
Diffuser
DS (ours)
Occlude first dim of velocity
43.3 2,2
59.0 : 2.3
68.7 +: 8.3
52.3 4 5.5
52.5 + 7.4
68.0 3,4
Occlude middle dim velocity
34.0
4.6
24.2 3.16 7.9
81.1 6,9
Ovclude last of velocity
51.1 1.3
10
61.6
32.7
70.7 4.7
110.9
Occlude first 2 dim of position
14.6 2.4
40.4
15.0
8.7
29.5 1.7
38.9 = 6.2
Average
3.7
69.4 2.8
5.9
313
44.2
74.7 4.2
D Ablation of reward modeling across MDP and POMDP.
We performed an ablation study on reward modeling, focusing on six locomotion environments. We action models: one set modeling the sequence of observations
trained two sets of transformer-bast and actions, and the other set incorporating the sequence 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 g emphasizing the benefits gained from incorporating reward modeling.
16

---

Table Ablation on reward modeling.
Action models
Enviromments
Transformer with reward modelling
Transformer without reward modelling
Hopper-medium-
76.6 14.2
69.7 13.4
MDP
Walker2d-medium-22
83.6 #:.3
82.7 #1.2
Halfcheetah -eiiretta-mmdiuu--22
47.8 L0.4
43,0 122.8
Average (MDP envs)
69.3
65.1
Hopper-medium--2
56.0 #3.5
43.6
POMDP
Walker2d-medium--22
74.3
54.2 #7.3
Halfihhetall-mcdiun-y
47.1
46.5 1.0.7
Averuge (POMDP cavs)
59.1
48.1
Average (All envs)
64.2
56.6
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, 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 two different orderings: Reward- Serrurd-Stttt-AAee State- Action (R, A) and State-Reware Rewa-de-eatioooo Aattonn (S, R, A). The results indicate that the S, R. A ordering outperforms the R, A ordering in the sulfcheetab-mediim- cedduummrepla-v2222 environment. DS models observations prior to rewards to be consistent with the functional definitions in MDPs, where the reward is typically function of observations (and potentially other variables), but not vice versa.
Ordering
haffcheccah-mccuun--rrlaaaaa Perferrr
Performance
R, S. A S, R
32.1 + 0.1
41.1 0.1
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-basse based approach to generate future plans. In this task with goal-conditioning, the position of the maze's 's goal serves the condition.
Regarding Diffuser, strictly adhere their codebase and observed that Diffuser impaints the last goal position into the difffusion sequence. Similarly, for DD, we follow their codebase and incorporate the goal as conditioning, embedding 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 serve as the goal position.
F11 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 the goal state. Consequently, when the agent close proximity to the goal point, the diffusion model plans a detour that initially takes the agent away from the goal and then redirects back towards the goal. Since the environment's reward function designed in such a way that the agent receives reward when it is near the goal, and the episode does
17

---

not terminate upon reaching the goal point but rather when it reaches the maximum episode length, it actually more advantgeoous for the agent remain 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 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 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 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 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
75
33
DOWFEC
Blange
Figure 4: Bar chart comparison of Diffuser, DD and DS with and without progressive conditioning in the single goal setting of Maze2D.
Multi Goal setting
133
Figure 5: Bar chart comparison of Diffuser, DD and DS with and without progressive conditioning in the multi goal setting of Maze2D.
18


Reviews:
[{"review_id": "hOYeCAT6gw", "reviewer": "Reviewer_Xfrm", "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.", "questions": "Please see weaknesses.", "limitations": "Yes.", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 2, "strengths": "1. The paper is well written and easy to read.\n2. 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.\n\n2. 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”. \n\n3. 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.\n\n4. Is there any intuition for why the conditioning in Equations (4), (5), (6) were chosen? \n\n5. 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your valuable comments and very detailed insights! We address your questions below.\n>Q1: Similar performance to DD on D4RL Gym tasks.\n\n1. This similarity, rather than reflecting a limitation of our method, highlights the near-saturated state of the D4RL benchmark, where many methods perform similarly. However, DS's top performance serves as a promising sanity check of its strength, even if the gains are close.\n\n2. Second, DS significantly surpasses other baseline methods in the POMDP setting, achieving a 76.9% improvement over DD and a 15.7% improvement over the next best method. During the rebuttal, a sensitivity analysis was conducted to consider various dimensions as detailed in Table 1 of the rebuttal PDF, showing that DS continues to outperform other baselines across different occlusions.\n\n3. In the Maze2D environment, the superior performance relative to baselines illustrates the strengths of DS's planning ability in complex single and multi-goal planning tasks.\n\n4. Apart from the above benchmark performance, we ran new experiments to demonstrate the compositional generalizability of DS which enables efficient transfer to new experiments, as detailed in Table 2 of the rebuttal PDF.\n\n>Q2: “Major similarities to previous work (trajectory transformer, decision diffuser) and the authors acknowledge this and highlight the reward modeling as novel… a (temporal difference) variant already exists in the trajectory transformer paper... the trajectory transformer with Q function (a form of reward modeling) guided planning which can be found in Section 4.2…”.\n\nWe respectfully disagree with both assertions:\n\n**Reward Modeling is Not the Only Novelty**: the main novelty of our work extends beyond reward modeling. This includes:\n\n1. DS introduces a novel token-first ordering of the trajectory components, instead of conventional timestep-first arrangements. This takes into account the distinct semantic attributes of each component, respecting their unique characteristics, such as domain type (discrete or continuous), and modality. While DT, TT and Diffuser chained states, action (and reward) with one model. Using a single model to model them together requires adjusting for modality-specific encoders and decoders, and close monitoring of training dynamics which can be avoided with DS.\n\n2. Modular Expressivity and Compositional Generalization: As in Eq. 3, each module is chained autoregressively with subsequent modules, leading to maximal expressivity under idealized conditions. The absence of parameter sharing across modules also enables hardware parallelization and easy transfer to new environments with shared modules. \n\n3. New experiments were conducted to showcase the compositional generalization of DS brought by the decomposed modules, as in Table 2 in the rebuttal PDF, where the same observation and/or reward models can be reused for different environments in Maze2D. \n\n4. Flexible Generative Parameterization: Our method allows the use of various deep generative models, catering to real-world environments where agents may face challenges such as executing discrete actions given continuous observations. \nSuperior Empirical Performance: The mentioned features translate into empirical superiority, as demonstrated in our results.\n\n**The TD Variant of TT cannot be used as justification to discount the novelty of our reward modeling approach**. While the Q-function in TT adopts a forward modeling approach, this necessitates dynamic programming. In contrast, DS takes a purely autoregressive route, signifying a clean and modular design. Unlike TT, which models cumulative returns, DS models per-time-step rewards. This differentiation is not trivial and translates into DS’s better empirical performance than TT.\n\n>Q3: explanation of which dimensions were removed and why is missing and a sensitivity analysis to the dimensions removed are missing as well.\n\nThanks for suggesting! We did a new sensitivity analysis on dimension occlusions, considering various scenarios and comparing them with all the baselines in the POMDP section. In the paper’s POMDPs experiment, two specific velocity dimensions of the joints were excluded from the full state representation for all the locomotion environments, simulating lacking relevant sensors. The results in Table 1 of the rebuttal PDF show that DS continues to outperform other baselines in aggregate performance.\n\n>Q4: any intuition for why the conditioning in Equations (4), (5), (6) were chosen?\n\nFrom 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, which is shown in Table 3 in the rebuttal PDF. Moreover, while not the focus of our experiments, our choices allows us to directly use pretrained perception models for mapping goals to observations. These models are often trained with data that does not include reward information,  e.g., text2video models a language command (goal) to a sequence of video frames (observation sequence).\n\n>Q5: any heuristic that could be used to choose the different modules with any online interaction? \n\nIn offline RL, datasets guide the choice of module architecture based on data type. For example, diffusion models may be chosen for continuous image-based states, while transformer models may be more suitable for less continuous actions or rewards. DS’s modular design reflects the modality and data types, which may not be appropriately addressed by employing a singular transformer or diffusion model for joint state-action modeling.\n"}, {"review_id": "86zMRcuU5P", "reviewer": "Reviewer_vCY4", "summary": "This 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. ", "questions": "1. 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.", "limitations": "The 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. ", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. 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.  ", "weaknesses": "In 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the valuable comments and very detailed insights! We address your questions below.\n>Q1: “...the training cost comparison to the other algorithms.”\n\nWe provide a comparison in Table 4 of the rebuttal PDF, showcasing the per-iteration cost. DS's training time is influenced by the choice of generative models. Since the state model has a high dimensionality, and the modules can be trained in parallel, the training time for DS is determined by the state model's training time. DS's training cost aligns with those of DT and TT when using a transformer state model. It is the same as DD's and on par with Diffuser's training cost for a diffusion state model. This result affirms that DS offers an efficient approach without compromising on training cost.\n\n>Q2: “... some of the claims are a bit strong.., 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\nWe apologize for the apparent confusion here. For expressivity, we indeed have a theoretical guarantee that follows directly from our autoregressive factorization in Eq. 3. That is, since autoregressive models are derived from chain rule of probability, under idealized assumptions on model expressivity and dataset size, they can learn any data distribution. For flexibility, our claim is informal and refers to the property that the modular design allows us to use any architecture and training objective for the observation, reward, and action modules. We will edit and clarify the language in the final version.\n\n>Q3: “The focus on MDP and POMDP environments was unclear… include some motivation or intuition as to why this approach would work better than previous works in both types…”\n\nBoth MDP and POMDP environments benefit from DS’s expressive modeling using the autoregressive factorization outlined in Eq. 3.\nFor MDPs, we can apply the D-separation rule to the corresponding graphical model (a simplification of Fig 2 with no observation nodes) to recognize that the next action is not d-separated from the immediate reward when conditioned on the current observation and the task goal. This insight motivates us to model rewards in DS (unlike DD), which proves beneficial as shown in our ablation study in Table 5.\nFor POMDPs, we can apply the D-separation rule to the corresponding graphical model (Fig 2) the D-separation rule to reveal that conditioning on the current observation and the goal does not obstruct the information pathway between the next action and prior or future observations, rewards, and actions, as demonstrated in Fig. 2. Thus, unlike prior works, the action and reward modules in DS by default condition on tokens from the past and future timesteps during training.\nFinally, DS's modular design respects the different modalities of states, actions, and rewards that typically span multiple dimensionalities and data types. This sets it apart from models like Diffuser, DT, and TT, which models them jointly with one architecture.\n\n>Q4: “Missing references [Mendez et al., 2022] and [Kaelbling, Littman, and Cassandra, 1998].”\n\nThanks! We will add these referenced papers in an updated version of our paper.\n\n>Q5: 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\nUni[MASK] falls into the transformer-based sequence modeling category, which we choose DT and TT for comparison. Uni[MASK] lacks reported results on locomotion tasks. Compared to Uni[MASK], DS also shows greater modularity. We see value in exploring comparisons with Uni[MASK] to deepen our understanding of training schemes.\n\n>Q6: 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\nIn Tables 1 and 2, values without reported errors are sourced from previous papers, while those with standard errors are from our experiments. We reported standard errors.\n\n>Q7: Could the authors please clarify what is meant by the expert-normalized scores in Section 4.2?\n\nThe term \"expert-normalized scores\" in the D4RL benchmark suite refers to a normalization scheme applied to the performance scores, where the upper and lower bounds are with respect to a fixed \"expert\" policy and a fixed \"random\" policy.\n\n>Q8: Why is the architectural flexibility investigated using Hopper-medium-v2 despite it not being used in the other experiments (e.g…in Table 3)?\n\nWe clarify that Hopper-medium-v2 was indeed used in the MDP and POMDP experiments, as in Tables 2 and 3. Choosing Hopper-medium-v2 for ablating architectural flexibility was arbitrary. The architectural flexibility experiment yields a 2x3x3 table of results where each entry is averaged over many random seeds, conducting on more environments would be computationally expensive.\n\n>Q9: Why does DS work better on medium but not medium expert or medium replay HalfCheetah (Table 3)?\n\nIn Table 3, DS and other methods exhibit a performance trend on Halfcheetah dataset, with the order: medium expert > medium > medium-replay. This suggests a hierarchy in data quality, where medium-replay contains more suboptimal trajectories than medium, and medium has more than medium-expert. \n>Q10: “How were the two dimensions of the state chosen to exclude (Section 4.3)? What is the observation representation…?”\n\nTwo velocity dimensions were excluded from the full state representation(consisting of positions and velocities) for POMDPs, simulating a lack of relevant sensors. A new sensitivity analysis on dimension occlusions further strengthens our results, as shown in Table 1 of the rebuttal PDF. DS continues to outperform other baselines.\n"}, {"review_id": "f74EVVHLl0", "reviewer": "Reviewer_E8pR", "summary": "This 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.", "questions": "I 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", "limitations": "Yes", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- 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.", "weaknesses": "- 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", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your valuable comments and very detailed insights! We address your questions below.\n\n>Q1: 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\nWe agree that new benchmarks can provide more interesting insights. We chose our benchmarks based on \n(a) those used in related works like DD, Diffuser, TT,  and we further extended them to expose challenging variations (eg, POMDPs), and \n(b) the compute available in our academic setting.\n\n1. We acknowledge that the D4RL benchmark is showing signs of saturation with many methods performing similarly. However, DS's top performance serves as a promising sanity check and indication of its relative strength.\n\n2. To create more challenging scenarios, we developed POMDP environments that occlue two velocity dimensions. During the rebuttal, we conducted a detailed analysis to consider many other variations. The results, compared with all baselines in the POMDP section and presented in Table 1 of the rebuttal PDF, show that DS continues to outperform other baselines across different ways of dimension occlusions.\n\n3. Lastly, our decision to use the Maze2D environment was done to evaluate the planning ability of DS. This is the key environment used for evaluating planning capabilities in our most closely related works, DD and Diffuser. The superior performance in this environment relative to baselines illustrates the strengths of our method in complex single and multi-goal planning tasks. \n\nWe acknowledge the potential value of exploring other environments such as RLBench and will explicitly list it for future work in the final version.\n\n>Q2: “... 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\nThanks for the interesting comment. 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, which is shown in Table 3 in the rebuttal PDF. Moreover, while not the focus of our experiments, our choices allows us to directly use pretrained perception models for mapping goals to observations. These models are often trained with data that does not include reward information,  e.g., text2video models a language command (goal) to a sequence of video frames (observation sequence).\n\n>Q3: “...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\nThank you for this insightful suggestion! We included some evidence in this regard with the multi-goal experiments in Maze2D, where an agent trained only on single goal trajectory transfers to multi-goal environments in a zero-shot manner at test-time. In line with the reviewer’s advice, we conducted experiments during the rebuttal period for new forms of generalization of our decomposed modules to new environments, as we describe next. \n\nSpecifically, we focused on the ability to transfer to other environments using shared observation models. We created modified versions of the Maze2D environment with varying action and reward spaces, while maintaining the same observation model. Our results demonstrate that DS facilitates an efficient learning process and compositional generalization. These findings further underline the modular efficiency of DS in scenarios where reusable components exist across different tasks or environments as shown in Table 2 in the rebuttal PDF.\n\n>Q4: 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\nWe do not believe marginalization is necessary for planning-driven execution of the policy., For any planning algorithm, we are interested in sampling from $p(a_t, o_{t+1:T}, r_{t+1:T} | G, o_{0:t}, a_{0:t−1},r_{t+1:T})$, akin to randomized shooting approaches. In the case of Decision Stacks, ancestral sampling, as described in Eqs. 6-8. If we were interested in density estimation for $p(a_t | G, o_{0:t}, a_{0:t−1},r_{t+1:T})$, then we would need to marginalize over future observations and rewards.\n\n>Q5: \"...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.\"\n\nThanks for the pointer and interesting idea! One potential execution strategy is to characterize poor reward distributions based on the extent to which they are out-of-distribution (ood)  from the training reward distribution. A highly ood sequence is likely to indicate poor planning of states, and we can exploit such a signal for replanning of states. While outside the scope of current work, we will include such a discussion in the paper for future investigation.\n"}, {"review_id": "uzpvEoFIfe", "reviewer": "Reviewer_xgrf", "summary": "This 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.", "questions": "(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\".", "limitations": "Methods 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.", "rating": 4, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "This 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.", "weaknesses": "One 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", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks for the valuable comments and very detailed insights! We address your questions below.\n>Q1: “One major concern regarding DS is the high complexity, as it requires training three generative models...the performance improvement compared to other baselines that use only one generative model is not significant.”\n\nWe respectfully disagree with the concerns about the complexity and performance of DS. On the contrary, we argue DS is simpler and more effective. \n\nComplexity:\n\n1. Each module is parameterized using the same kind of architecture (transformers or diffusion model) as the other baselines. Training time for DS doesn't increase compared to the close baseline DD, despite using three generative models. This is because DS's modularity allows for the parallel training of the three modules. The training time is mainly influenced by the state model of DS due to its high dimensionality. A training cost comparison in Table 4 of the rebuttal PDF confirms that DS's training cost aligns with those of DT, TT, DD, and Diffuser.\n\n2. Enhanced modality-specific flexibility: Although baselines like DT and TT which use a single generative model, bring surface-level simplicity, we believe DS's modular design is less complex as it explicitly respects the different modalities for states, actions, rewards. For example, actions can be less continuous than the states. Using a single diffusion or transformer model to model them together requires adjusting for modality-specific encoders and decoders, as well as close monitoring of training dynamics.\n\n\nPerformance Gain:\n\n1. In the MDP setting where benchmarks are nearly saturated as recent methods such as DD and TT achieve similar results, DS still performs better, which serves as an essential sanity check and indicates its strong empirical performance, which we validated with our subsequent experiments.\n\n2. DS demonstrates robust superiority in more complex environments, such as the POMDP setting and the Maze2D environment. The improvements over other methods in these settings are more substantial, indicating DS’s expressive modeling derived from chain rules and D-separation rule which can be applied to the graphical model in Fig 2 translate into tangible performance gains in more complex tasks.\n\n3. DS modular design allows efficient compositional generalization. We demonstrate this on maze2D where modules can be reusable across envs, as shown in Table 2 of the rebuttal PDF.\n\n\n>Q2: lacks an explanation of the specific dimensions that were excluded in the POMDP setting. \n\nIn Section 4.3, we excluded two specific velocity dimensions of the joints from the full state representation for all locomotion environments, simulating lacking relevant sensors. We appreciate your inquiry into this design choice and conducted a sensitivity analysis on dimension occlusions in Table 1 of the rebuttal PDF. The new analysis considered various exclusions and the results affirm that DS continues to perform robustly, outperforming other baselines.\n\n> Q3: “... how training can be parallelized and how the choice of generative models affects performance.”\n\nDS modules’ parallel training is achieved by a technique known as teacher forcing, where during training, the true previous output is fed as input to the next time step. We leverage the computational efficiency of parallel processing without losing the integrity of the dependencies among the models.\n\n>Q4: “ 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 in Table 4? … provide corresponding reasonable explanations instead of simple combination attempts.” “... lacks rational analysis in the selection of generative models to generate higher quality trajectories. “\n\nIn our experiment, we chose models aligned with the properties of the data modality: diffusion models for the continuous state sequences, and transformers for the more high-frequency action and reward sequences. This wasn't a combination attempt. For action sequences, which often appear as joint torques and tend to be more high-frequency and less smooth, we selected the transformer model. The same applies to our choice of the reward model. This decision was based on the transformer's suitability for handling such high-frequency and varied data structures, as evidenced by its superior performance in Table 4 (55.2 vs. 53.5 for the diffusion action model). By aligning the models with the nature of the data, we optimized our architecture to reflect the unique characteristics of the modalities in RL environments.\n\n>Q5: \"...handcoded controller implemented in 4.1 and 4.2? ... role? 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\nThe handcoded controller is an existing technique used in Diffuser, which calculates action based on planned waypoints with the physics dynamics. It computes the target distance and optimizes actions. Though DD and DS perform similarly with this controller, the key distinction lies in executing plans with action models. Diffuser overlooks the varied data nature between states and actions by diffusing them together, while DD uses an MLP dynamics model that does not consider sufficient contexts. DS employs a more expressive approach with more contextual information. While planning in DD and DS may appear similar, the execution varies, and DS shows superior performance due to its modular structure, especially in the nuanced execution phase. The modular structure of DS does indeed have a significant impact, specifically in the more nuanced execution phase, as opposed to merely selecting a good generative model.\n"}, {"review_id": "8oDTuF8DrG", "reviewer": "Reviewer_ywNv", "summary": "This 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.", "questions": "- 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?", "limitations": "- 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. ", "rating": 3, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "- The paper is fairly well presented and easy to follow. \n\n- The proposed method is extensively evaluated against multiple related approaches.", "weaknesses": "- 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks for the valuable comments and very detailed insights! We address your questions below.\n\n> Q1: “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\nWe hope the following clarify DS's novel modular and expressive design:\n\n1. The performance on D4RL Mujoco MDPs are saturated with a number of recent methods performing quite close. While our gains are close, the fact that DS is at the top serves as a promising sanity check and indication of its strong performance. \n\n2. We developed more challenging POMDP environment by excluding two velocity dimensions. During this rebuttal period, new experiments were conducted for dimension occlusion sensitivity analysis. The results, shown in Table 1 in the rebuttal PDF, strengthens DS's superior performance in POMDP cases. \n\n3. On the Maze2D environment, we showcase the planning ability of DS. This is the key environment used for evaluating planning capabilities in our most closely related works, DD and Diffuser. The superior performance in this environment relative to baselines illustrates the strengths of our method in complex single and multi-goal planning tasks. \n\n4. New experiments were conducted to showcase the compositional generalization of DS brought by the decomposed modules, as in Table 2 in the rebuttal PDF, where the same observation and/or reward models can be reused for different environments in Maze2D. \n\n5. DS's modular design is less complex as it explicitly respects the different modalities for states, actions, rewards. For example, actions can be less continuous than the states. Using a single diffusion or transformer model to model them together requires adjusting for modality-specific encoders and decoders, as well as close monitoring of training dynamics.\n\n\n> Q2: “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.” “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\nDS introduces a novel token-first ordering of the trajectory components instead of the timestep-first arrangements. This design, while appearing counter-intuitive, serves 2 important functions: \n\n(1)  A modality-focused approach explicitly emphasizes the role of planning. Temporal ordered models such as DT do not permit planning. \n\n(2) It caters to the distinct semantic attributes of observations, actions, and rewards. It recognizes the unique characteristics such as dimensionality and domain type (discrete or continuous). By allowing for separate modeling, DS provides flexibility in architectural choices to cater to these differences, unlike other methods such as DT, TT and Diffuser that chain states, actions, (and rewards) together. The promising results in our experiments, especially on the POMDPs and Maze2D have validated DS with superior expressivity and planning ability. This modular design also brings novel compositional generalizations.\n\n>Q3: “Lack of evidences to support the claims on modular expressivity (L145)...any experiments showing it can transfer to new environments?”  “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\nIn line with your suggestion, we conducted new experiments to explore the potential of our decomposed submodules for compositional generalization and efficient multitask learning. Specifically, we focus on the ability to transfer to other environments using shared observation models. We created modified versions of the Maze2D environment with varying action and reward spaces (details in Table 2 caption), maintaining the same observation model. Our results demonstrate that DS facilitates an efficient compositional generalization in scenarios where reusable components exist across different environments as detailed in Table 2 in the rebuttal PDF.\n\n>Q4: “The empirical results, especially those on D4RL locomotion tasks, are not significant enough to justify the extra compute introduced.”\n\n1. We justified the significance of our results in Q1. \n\n2. No significant extra compute: Each module is parameterized using the same kind of architecture (transformers or diffusion) as the other baselines. The training time for DS does not increase as compared to the closest baseline DD despite having three generative models, due to DS's parallel module training. The training time is determined by the state model of DS due to high dimensionality. A training cost comparison in Table 4 of the rebuttal PDF confirms that DS's training cost aligns with other baselines.\n\n>Q5: Why DS's performance on POMDP Hopper task with medium-expert data is even better than that on fully observable setting?\n\nThe performances in POMDP and MDP settings are similar within the error bar range. The higher performance in the POMDP might stem from the limited 15 rollouts of a random seed. This suggests that DS is capable of extracting enough useful information to infer hidden states. This trend is not unique to DS, e.g., DT's performance on the Hopper-Medium-Expert task is close in both the POMDP and MDP cases.\n\n>Q6: How does DS compare to more straightforward baseline, such as the one with separate heads for three modules but shares the same backbone?\n\nThis baseline resembles the Trajectory Transformer, which predicts the states, actions, and rewards with a common backbone architecture. In comparison, DS attains performance gains of 2.7% on MDPs and 40.2% on POMDPs over TT. Further, compared to such baselines, DS's decomposed modules bring more architectural flexibility (Table 4) and transferrability across environments (Table 2, rebuttal pdf). \n"}]

Author rebuttal:
We thank all the reviewers for providing insightful feedback and constructive suggestions, including recommendations for new ablation studies. In alignment with these insightful suggestions, we have conducted and provided 4 new experimental results that clarify Decision Stacks's novelty and address the concerns raised in the reviews. 

The summary of these experiments is as follows:

1. Table 1: **Sensitivity Analysis on Occluded Dimensions for POMDPs**. In the Hopper environment, the full state encompasses 5 dimensions of position data and 6 dimensions of velocity data across different joints. The table outlines experiments involving various dimensions and semantics for occlusion. DS consistently exhibits superior or second-best performance compared to other baselines on the hopper-medium-expert-v2 dataset.

     Results: **DS continues to outperform other baselines across different ways of dimension occlusions in aggregate performance.**



2. Table 2: **Experiments on Compositional Generalization in Maze2D**. With the absence of parameter sharing, the DS modular design allows efficient compositional generalization. This experiment investigates the transfer to other environments using shared observation models with diverse action and reward spaces. In the Maze2D environment, both dense and sparse rewards are analyzed, along with three variations of force applied to a 2D ball. These results emphasize DS's modular efficiency across scenarios where reusable components exist across different tasks or environments. 

   Results: **DS modular design allows efficient compositional generalization**

3. Table 3: **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 states, rewards, and actions. Specifically, we ordered states 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 and found our choice to significantly outperform the counterpart. This table compares two different orderings: Reward-State-Action (R, S, A) and State-Reward-Action (S, R, A). The findings reveal that the S, R, A ordering outperforms the R, S, A ordering in the halfcheetah-medium-replay-v2 environment, underscoring DS's consistency with functional definitions in MDPs.


   Results: **DS's choice of S, R, A ordering outperforms the R, S, A ordering**


4. Table 4: **Per-Iteration Training Cost for Different Algorithms**. We provide a comparison of the per-iteration cost. DS's training time is influenced by the choice of generative models. Since the state model has a high dimensionality, and the modules can be trained in parallel, the training time for DS is determined by the state model's training time. Specifically, DS's training cost aligns with those of DT and TT (transformer-based) and is consistent with DD's and on par with Diffuser's (diffusion-based) training cost. This result further affirms that DS delivers an efficient approach without compromising on training cost.

   Results: **DS's training cost aligns with other baselines.**

Please refer to the attached PDF for details of the experiments

Meta-review: The reviews and discussion offered mixed reactions and support. However, the clarity of presentation and analysis of the approach is relatively well done, and the modular design complements existing approaches. We do encourage the authors to use the feedback to revise and strengthen this work in the coming weeks.
Decision: Accept (poster)