SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Siyan Zhao
Department of Computer Science  
University of California Los Angeles  
siyanz@cs.ucla.edu
Department of Computer Science  
University of California Los Angeles  
adityag@cs.ucla.edu
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.<sup>1</sup>
Modularity is a critical design principle for both software systems and artificial intelligence (AI). It allows for the creation of flexible and maintainable systems by breaking them down into smaller, independent components that can be easily composed and adapted to different contexts. For modern deep learning systems, modules are often defined with respect to their input and output modalities and their task functionalities. For example, Visual ChatGPT [Wu et al., 2023] defines a family of  $22+$  vision and language foundation models, such as ChatGPT [OpenAI, 2022] (language generation), SAM [Kirillov et al., 2023] (image segmentation), and StableDiffusion [Rombach et al., 2022] (text-to-image generation) for holistic reasoning over text and images. In addition to enabling new compositional applications, modularity offers the promise of interpretability, reusability, and debugging for complex workflows, each of which poses a major challenge for real-world AI deployments.
This paper presents progress towards scalable and flexible reinforcement learning (RL) through the introduction of a new modular probabilistic framework based on deep generative models. Prior work in modular RL focuses on spatiotemporal abstractions that simplify complex goals via hierarchical RL, e.g., [McGovern and Barto, 2001, Andreas et al., 2017, Simpkins and Isbell, 2019, Ahn et al., 2022, Kulkarni et al., 2016, Mendez et al., 2021]. Distinct but complementary to the prior lines of work, our motivating notion of modularity is based on enforcing token-level hierarchies in generative models of trajectories. In the context of RL, trajectories typically consist of a multitude of different tokens of information: goals, observations, rewards, actions. As shown in many recent works [Chen et al., 2021, Janner et al., 2021, 2022, Ajay et al., 2022, Zheng et al., 2022, Reed et al., 2022], we can effectively reduce RL to probabilistic inference [Levine, 2018] via learning deep generative models
<sup>1</sup>The project website and code can be found here: https://siyan-zhao.github.io/decision-stacks/
over token sequences. However, these frameworks lack any modular hierarchies over the different tokens leading to adhoc choices of generative architectures and objectives, as well as conditional independence assumptions that can be suboptimal for modeling long trajectory sequences.
We introduce Decision Stacks, a family of generative algorithms for goal-conditioned RL featuring a novel modular design. In Decision Stacks, we parameterize a distinct generative model-based module for future observation prediction, reward estimation, and action generation and chain the outputs of each module autoregressively. See Figure 1 for an illustration. While our factorization breaks the canonical time-induced causal ordering of tokens, we emphasize that the relative differences in different token types is significant to necessitate token-level modularity for learning effective policies and planners. Besides semantic differences, the different token types also show structural differences with respect to dimensionalities, domains types (discrete or continuous), modalities (e.g., visual observations, numeric rewards), and information density (e.g., rewards can be sparse, state sequences show relatively high continuity). Instead of modeling the token sequence temporally, parameterizing a distinct module for each token type can better respect these structural differences. In contrast, previous works like Decision Transformer Chen et al. [2021], Trajectory Transformer Janner et al. [2021] and Diffuser Janner et al. [2022] chained states, actions, and, in some cases, rewards within a single temporal model.
In practice, we can train each module in Decision Stacks independently using teacher forcing [Williams and Zipser, 1989], which avoids additional training time as the three modules can be trained in parallel. Decision Stacks shares similarities with many recent works [Janner et al., 2021, Ajay et al., 2022, Janner et al., 2022] that aim to reduce planning to sampling from a generative model. However, our modular design offers additional flexibility and expressivity. Each generative module itself is not restricted to being an autoregressive model and we experiment with modules based on transformers, diffusion models, and novel hybrids. Each generative modeling family makes tradeoffs in architecture, sampling efficiency, and can show varied efficacy for different data modalities. Decision Stacks also brings compositional generalization in scenarios where the modules can be reused across varied tasks or environments. A modular design that easily allows for the use of arbitrary generative models, along with an autoregressive chaining across the modules permits both flexibility and expressivity.
Empirically, we evaluate Decision Stacks on a range of domains in goal-conditioned planning and offline RL benchmarks for both MDPs and POMDPs. We find that the joint effect of modular expressivity and flexible parameterization in our models provides significant improvements over existing offline RL methods. This holds especially in partially observable settings, where Decision Stacks achieves a  $15.7\%$  performance improvement over the closest baseline, averaged over 9 offline RL setups. We also demonstrate the flexibility of our framework by extensive ablation studies over the choice of generative architectures and inputs for each module.
2.1 Goal conditioned POMDPs
We operate in the formalism of goal-conditioned Partially Observable Markov Decision Processes (POMDP) defined by the tuple  $\mathcal{M} \coloneqq (\mathcal{O},\mathcal{S},\mathcal{A},\mathcal{G},\mathcal{P},\mathcal{R},\mathcal{E},\gamma ,p_0(s),T)$ . Respectively,  $\mathcal{O}$  and  $\mathcal{S}$  denote the observation space and the underlying state space, which are fully observable in the case of MDPs.  $\mathcal{A}$ , the action space, is consistent with that of MDPs. In the goal-condition setting,  $\mathcal{G}$  specifies the task goal distribution which could be e.g., a language instruction or a visual destination state for multi-task policies, or a designed cumulative return for single-task policies. The transition probability function,  $\mathcal{P}:S\times \mathcal{A}\times S\to [0,1]$ , describes the transition dynamics. Meanwhile,  $\mathcal{R}:G\times S\times A\mapsto \mathbb{R}$  defines the rewards that the decision-maker receives after performing action  $a$  in state  $s$ . The observation emission model  $\mathcal{E} = P(o|s)$ , determines the probability of observing  $o$  in state  $s$ . The  $\gamma$ ,  $p_0(s_0)$ , and  $T$  denote the discount factor [Puterman, 2014], initial latent state distribution, and horizon of an episode. In a POMDP context, the observations generated from the underlying state are intrinsically non-Markovian. The goal-conditioned RL objective is to find the optimal policy  $\pi^*$  that maximizes the expected cumulative discounted reward over the episode horizon:  $\eta_{\mathcal{M}} \coloneqq \mathbb{E}_{G\sim \mathcal{G},a_t\sim \pi (\cdot |s_t,g),s_{t + 1}\sim \mathcal{P}(\cdot |s_t,a_t)}\left[\sum_{t = 0}^{T}\gamma^t r(s_t,a_t,G)\right]$ .
Figure 1: 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 (RL) is a paradigm for policy optimization where the agent is only given access to a fixed dataset of trajectories and cannot interact with the environment to gather additional samples. Offline RL can be useful in domains where collecting data online is challenging or infeasible such as healthcare [Murphy et al., 2001] and autonomous driving. A major obstacle in offline RL is dealing with distributional shifts. If we naively use Bellman backup for learning the Q-function of a given policy, the update which relies on the actions sampled from policy  $\pi$  can learn inaccurately high values for out-of-distribution actions, leading to instability in the bootstrapping process and causing value overestimation [Kumar et al., 2020].
2.3 Generative models: Autoregressive transformers and Diffusion models
In this work, we are interested in learning the distribution  $p_{\mathrm{data}}(\mathbf{x}|\mathbf{c})$  using a dataset  $D$  consisting of trajectory samples  $\mathbf{x}$  and conditioning  $\mathbf{c}$ . We consider two conditional generative models for parameterizing our agent policies to learn the distribution:
Transformer is a powerful neural net architecture for modeling sequences [Vaswani et al., 2017]. It consists of multiple identical blocks of multi-head self-attention modules and position-wise fully-connected networks. The vanilla transformer can be modified with a causal self-attention mask to parameterize an autoregressive generative model as in GPT [Radford et al., 2018]. Autoregressive generative models, such as the transformers, factorize the joint distribution  $p(x_{1},\ldots ,x_{n})$  as a product of conditionals, which can be represented as:  $p(\mathbf{x}) = \prod_{i = 1}^{n}p(x_{i}|\mathbf{x}_{< i})$ . This equation shows that the probability of each variable  $x_{i}$  depends on the previous variables  $x_{1},\ldots ,x_{i - 1}$ . One advantage of this factorization is that each conditional probability can be trained independently in parallel via teacher forcing [Williams and Zipser, 1989]. In autoregressive generation, sampling is done sequentially, where each variable  $x_{i}$  is sampled based on its preceding variables.
Diffusion Models [Sohl-Dickstein et al., 2015, Ho et al., 2020] are latent variable models that consist of a predefined forward noising process  $q(\pmb{x}_{k+1}|\pmb{x}_k) \coloneqq \mathcal{N}(\pmb{x}_{k+1}; \sqrt{\alpha_k}\pmb{x}_k, (1 - \alpha_k)I)$  that gradually corrupts the data distribution  $q(\pmb{x}_0)$  into  $\mathcal{N}(0, I)$  in  $K$  steps, and a learnable reverse denoising process  $p_{\theta}(\pmb{x}_{k-1}|\pmb{x}_k) \coloneqq \mathcal{N}(\pmb{x}_{k-1}|\mu_{\theta}(\pmb{x}_k, k), \Sigma_k)$ . For sampling, we first generate a latent sample from the Gaussian prior  $\mathcal{N}(0, I)$  and gradually denoise it using the learned model  $p_{\theta}(\pmb{x}_{k-1}|\pmb{x}_k)$  for
Figure 2: Graphical model for the data generation process in a POMDP. Here, we show the case where the behavioral policy can potentially act based on hidden state information. The dashed circles imply that this state information is not stored in the offline dataset.  $G$  represents the task conditioning, e.g., a target return (for single-task agents) or a navigation goal (for multi-task agents).
$K$  steps to obtain the data sample  $\pmb{x}_0$ . Diffusion models can be extended for conditional generation using classifier-free guidance [Ho and Salimans, 2022] where a conditional  $\epsilon_{\theta}(\pmb{x}_k,\mathbf{c},k)$  and an unconditional  $\epsilon_{\theta}(\pmb{x}_k,k)$  noise model is trained. Conditional data is sampled with the perturbed noise  $\epsilon_{\theta}(\pmb{x}_k,k) + \omega (\epsilon_{\theta}(\pmb{x}_k,\mathbf{c},k) - \epsilon_{\theta}(\pmb{x}_k,k))$ , where  $\omega$  is the guidance strength and  $\mathbf{c}$  is the conditioning information.
3 Flexible and Modular RL via Decision Stacks
In Figure 2, we consider a directed graphical model for the data generation process in Partially Observable Markov Decision Processes (POMDP) Kaelbling et al. [1998]. The environment encodes the underlying state transitions  $P(s_{t + 1}|s_t,a_t)$ , goal-dependent reward function  $P(r_{t}|s_{t},a_{t},G)$ , and the observation emission probability  $P(o_{t}|s_{t})$ . Unlike the learning agent which only has access to observations, the behavioral policy used for generating the offline trajectory dataset might also have access to the hidden state information. Such a situation is common in real-world applications, e.g., a human demonstrator may have access to more information about its internal state than what is recorded in video demonstration datasets available for offline learning. Finally, note that a Markov Decision Process (MDP) can be viewed as a special case of a POMDP, where the observation at each timestep,  $o_t$ , matches the underlying state,  $s_t$ . To avoid ambiguity, we overload the use of  $o_t$  to denote both the state and observation in an MDP at time  $t$ .
In the context of goal-conditioned decision-making, a finite-horizon trajectory in the offline dataset  $\mathcal{D}$  is composed of a goal  $G$  and a sequence of observations, actions, and reward tokens.
$$
\tau = \left(G, o _ {0}, a _ {0}, r _ {0}, \dots , o _ {t}, a _ {t}, r _ {t}, \dots , o _ {T}, a _ {T}, r _ {T}\right). \tag {1}
$$
Our primary objective lies in learning a goal-conditioned distribution for  $P_{\mathrm{data}}(a_{0:T}, o_{1:T}, r_{0:T} | o_0, G)$  conditioned on an arbitrary goal  $G \in \mathcal{G}$  and an initial observation  $o_0 \in \mathcal{O}$ . Leveraging the chain rule of probability, we can factorize this joint distribution into a product of conditional probabilities. For example, Janner et al. [2021] use a time-induced autoregressive factorization:
$$
P _ {\text {d a t a}} \left(a _ {0: T}, o _ {1: T}, r _ {0: T} \mid o _ {0}, G\right) \approx \prod_ {t = 1} ^ {T} P _ {\theta} \left(o _ {t} \mid \tau_ {<   t}\right) \prod_ {t = 0} ^ {T} P _ {\theta} \left(r _ {t} \mid o _ {t}, \tau_ {<   t}\right) \prod_ {t = 0} ^ {T} P _ {\theta} \left(a _ {t} \mid o _ {t}, r _ {t}, \tau_ {<   t}\right) \tag {2}
$$
where  $\tau_{<t}$  denotes all the tokens in the trajectory before time  $t$ . Each conditional factor is parameterized via an autoregressive transformer with shared parameters  $\theta$ . If the parameterization is sufficiently expressive, any choice of ordering for the variables suffices. However, in practice, we are limited by the size of our offline dataset and the choice of factorization can play a critical role.
In Decision Stacks, we propose to use a modular factorization given as:
$$
P _ {\text {d a t a}} \left(a _ {0: T}, o _ {1: T}, r _ {0: T} \mid o _ {0}, G\right) \approx \underbrace {P _ {\theta_ {1}} \left(o _ {1 : T} \mid o _ {0} , G\right)} _ {\text {o b s e r v a t i o n m o d u l e}} \cdot \underbrace {P _ {\theta_ {2}} \left(r _ {0 : T} \mid o _ {0 : T} , G\right)} _ {\text {r e w a r d m o d u l e}} \cdot \underbrace {P _ {\theta_ {3}} \left(a _ {0 : T} \mid o _ {0 : T} , r _ {0 : T} , G\right)} _ {\text {a c t i o n m o d u l e}}. \tag {3}
$$
Each of the 3 modules (observations, rewards, or actions) focuses on predicting a distinct component of the POMDP and has its own set of parameters  $(\theta_{1},\theta_{2},\theta_{3})$ . Our motivation stems from the fact that in real-world domains, each component is sufficiently distinct from the others in its semantics and representation. Such variances span across a multitude of factors including dimensionalities, domain types (discrete or continuous), modalities (e.g., visual observations, numeric rewards), and information density (e.g., rewards can be sparse, state sequences show relatively high continuity).
Modular Expressivity. In Eq. 3, each module is chained autoregressively with the subsequent modules. This is evident as the output variables of one module are part of the input variables for all the subsequent modules. Under idealized conditions where we can match each module to the data conditional, this autoregressive structure enables maximal expressivity, as autoregressive models derived from the chain rule of probability can in principle learn any data distribution given enough model capacity and training data. Further, our explicit decision to avoid any parameter sharing across modules also permits trivial hardware parallelization and transfer to new environments with shared structure.
Flexible Generative Parameterization. Since each module predicts a sequence of objects, we use any deep generative model for expressive parameterization of each module. Our experiments primarily focus on autoregressive transformers and diffusion models. We also consider hybrid combinations, as they are easy to execute within our framework and can avoid scenarios where individual model families suffer, e.g., diffusion models lag behind transformer for discrete data; whereas transformers are generally poor for modeling continuous signals such as image observations. In real-world environments, many of these challenges could simultaneously occur such as agents executing discrete actions given continuous observations. Finally, each module is conditioned on a goal. For training a multi-task agent, the goal can be specified flexibly as spatial coordinates, a visual image, a language instruction, etc. For single-task agents, we specify the goal as the trajectory returns during training and desired expert-valued return during testing, following prior works in return-conditioned offline RL [Chen et al., 2021, Emmons et al., 2021].
Learning and Inference. Given an offline dataset, each module can be trained in parallel using teacher forcing [Williams and Zipser, 1989]. At test-time, our framework naturally induces a planner, as in order to predict an action at time  $t$ , we also need to predict the future observations and rewards. We can execute either an open-loop or closed-loop plan. Open-loop plans are computationally efficient as they predict all future observations and rewards at once, and execute the entire sequence of actions. In contrast, a closed-loop plan is likely to be more accurate as it updates the inputs to the modules based on the environment outputs at each time-step. Using a closed-loop plan, we can sample the action at time  $t$  as follows:
$$
\hat {o} _ {t + 1: T} \sim P _ {\theta_ {1}} \left(o _ {t + 1: T} \mid o _ {0: t}, G\right) \tag {4}
$$
$$
\hat {r} _ {t + 1: T} \sim P _ {\theta_ {2}} \left(r _ {t + 1: T} \mid r _ {0: t}, o _ {0: t}, \hat {o} _ {t + 1: T}, G\right) \tag {5}
$$
$$
\hat {a} _ {t} \sim P _ {\theta_ {3}} \left(a _ {t} \mid a _ {0: t - 1}, o _ {0: t}, \hat {o} _ {t + 1: T}, r _ {0: t}, \hat {r} _ {t + 1: T}, G\right) \tag {6}
$$
The hat symbol  $(\wedge)$  indicates predicted observations, rewards, and actions, while its absence denotes observations, rewards, and actions recorded from the environment and the agent in the previous past timesteps. For closed-loop planning, Eqs. 4, 5, 6 require us to condition the joint observation, reward and action distributions on the past trajectory tokens. For a module that is parameterized autoregressively, this is trivial as we can simply choose a time-induced ordering and multiply the conditionals for the current and future timesteps. For example, if the observation module is an autoregressive transfer, then we can obtain the sampling distribution in Eq. 4 as:  $P_{\theta_1}(o_{t + 1:T} \mid o_{0:t}, G) = \prod_{i = t + 1}^{T} P_{\theta_1}(o_i \mid o_{<i}, G)$ . For a diffusion model, this task is equivalent to inpainting and can be done by fixing the environment observations until time  $t$  at each step of the denoising process [Janner et al., 2022].
Distinction with Key Prior Works. We will include a more detailed discussion of broad prior works in §5 but discuss and contrast some key baselines here. While the use of generative models for goal-conditioned offline RL is not new, there are key differences between Decision Stacks and recent prior works. First, we choose a planning approach unlike other model-free works, such as Decision Transformers [Chen et al., 2021] and diffusion-based extensions [Wang et al., 2022]. Second, there exist model-based approaches but make different design choices; Trajectory Transformer [Janner et al., 2021] uses a time-induced causal factorization parameterized by a single autoregressive transformer, Diffuser [Janner et al., 2022] uses diffusion models over stacked state and action pairs, Decision Diffuser [Ajay et al., 2022] uses diffusion models for future state prediction and an MLP-based inverse dynamics model to extract actions. Unlike these works, we propose a modular structure that is maximally expressive as it additionally models reward information and does not make any conditional independence assumption for the state, reward and action modules. As our experiments demonstrate, the modular expressivity and architectural flexibility in Decision Stacks are especially critical for goal-conditioned planning and dealing with partial observability.
Table 1: Performance on Maze2D tasks. DS significantly outperforms other baselines without the need for a handcoded controller. Note that DD and DS share the same diffusion-based 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 ( $\geq 0.95 \cdot \max$ ).
Task	Environment	MPPI	CQL	IQL	Diffuser	DD	DS	Diffuser with Handcoded Controller	DS / DD with Handcoded Controller
Single Goal	umaze	33.2	5.7	47.4	86.9 ±26.4	113.8 ±11.3	111.3 ±12.2	113.9 ±3.1	119.5 ±2.6
medium	10.2	5.0	34.9	108.5 ±17.4	103.7 ±21.2	111.7 ±2.4	121.5 ±2.7	112.9 ±11.8
large	5.1	12.5	58.6	45.4 ±14.5	111.8 ±43.4	171.6 ±13.4	123.0 ±6.4	132.8 ±21.0
Average	16.2	7.7	47.0	80.2	109.8	131.5	119.5	121.7
Multi Goals	umaze	41.2	-	24.8	114.4 ±16.3	105.6 ±14.5	121.3 ±12.2	129.0 ±1.8	136.1 ±4.2
medium	15.4	-	12.1	54.6 ±14.5	126.4 ±14.3	122.3 ±3.7	127.2 ±3.4	124.6 ±11.3
large	8.0	-	13.9	41.0 ±20.1	116.0 ±33.1	126.7 ±21.8	132.1 ±5.8	134.8 ±12.3
Average	21.5	-	16.9	70.0	111.6	123.4	129.4	131.8
Our experiments aim to answer the following questions:
$\S 4.1$  How does Decision Stacks perform for long-horizon multi-task planning problems?
$\S 4.2$  How does Decision Stacks compare with other offline RL methods in MDP environments?
$\S 4.3$  How does Decision Stacks compare with other offline RL methods in POMDP environments?
$\S 4.4$  How does the architectural feasibility for each module affect downstream performance? How does the modularity enable compositional generalization? How important is the role of reward modeling for Decision Stacks?
For §4.1, §4.2, and §4.3, we experiment with D4RL environments and parameterize Decision Stacks with a diffusion-based observation model, an autoregressive transformer-based reward model, and an autoregressive transformer-based action model. Finally, in §4.4, we will ablate the full spectrum of architecture design choices for each module.
4.1 Long-Horizon Goal-Conditioned Environments