SlowGuess/NeurIPS_2023_OCR huggingface dataset:
We first test for the planning capabilities of Decision Stacks on the Maze2D task from the D4RL [Fu et al., 2020] benchmark. This is a challenging environment requiring an agent to generate a plan from a start location to a goal location. The demonstrations contain a sparse reward signal of  $+1$  only when the agent reaches close to the goal. Following Janner et al. [2022], we consider 2 settings. In the Single Goal setting, the goal coordinates are fixed, and in the Multi Goal setting, the goals are randomized at test-time. We compare against classic trajectory optimization techniques that have knowledge of the environment dynamics (MPPI [Williams et al., 2015], extensions of model-free RL baselines (CQL [Kumar et al., 2020] and IQL [Kostrikov et al., 2021]), and the two most closely related works in generative planning based on diffusion models: Diffuser [Janner et al., 2022] and Decision Diffuser (DD) [Ajay et al., 2022].
We show our results in Table 1 for different goal types and maze grids. While Janner et al. [2022] previously demonstrated remarkable ability in generating long-horizon plans using Diffuser, their trajectory plans were executed by a handcoded controller. However, we experimentally found that Diffuser and DD's own generated actions fail to perfectly align with their generated plans, as shown in the example rollouts in Figure 3. We hypothesize this could stem from the lack of modularity in Diffuser affecting the generation fidelity, or the lack of expressivity in using an MLP-based inverse dynamics model in DD which limits the context length required for long-horizon planning. In contrast, we find that DS generates robust trajectory plans and matching action sequences with significant improvements over baselines.
4.2 Offline Reinforcement Learning Performance in MDPs
Next, we examine the performance of Decision Stacks in offline RL tasks across various 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 et al., 2021] and CQL [Kumar et al., 2020], model-based algorithm MOReL [Kidambi et al., 2020], transformer-based
Figure 3: Example rollouts on the Maze2D-medium-v1 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 demonstrate that Diffuser, DD, and DS are all capable of generating good plans that can be executed well with a handcoded controller. However, the respective action models result in differing executions. Compared to DD and Diffuser, DS generates smoother trajectories that are more closely aligned with the future waypoints planned by the observation model.
generative models such as Decision Transformer (DT) [Chen et al., 2021] and Trajectory Transformer (TT) [Janner et al., 2021], and diffusion-based generative models Diffuser [Janner et al., 2022] and Decision Diffuser (DD) [Ajay et al., 2022]. In our evaluation, we also included our reproduced scores for DD. DD uses the same architecture for observation prediction as Decision Stacks and is hence, the closest baseline. However, we found its performance to be sensitive to return conditioning and in spite of an extensive search for hyperparameters and communication with the authors, our reproduced numbers are slightly lower. We provide more details in the Appendix. For a fair comparison, we used the same set of hyperparameters that give the best performance for the DD baseline.
We show results averaged over 15 planning seeds and normalize the scores such that a value of 100 represents an expert policy, following standard convention [Fu et al., 2020]. Decision Stacks outperforms or is competitive with the other baselines on 6/9 environments and is among the highest in terms of aggregate scores. These results suggest that even in environments where we can make appropriate conditional independence assumptions using the MDP framework, the expressivity in the various modules of Decision Stacks is helpful for test-time generalization.
4.3 Offline Reinforcement Learning Performance in POMDPs
Next we consider the POMDP setting where the logged observations are incomplete representations of the underlying states. To generate the POMDPs datasets, we exclude the two velocity dimensions from the full state representation, which consists of both positions and velocities. This simulates a lack of relevant sensors. A sensitivity analysis on dimension occlusions further strengthens our results, as shown in Appendix Table 7. DS continues to outperform other baselines for each environment from the D4RL locomotion datasets. We report our results in Table 3 and compare against other generative baselines. Decision Stacks (DS), consistently achieves competitive or superior results compared to the other algorithms, including BC, DT, TT, and DD. Notably, DS outperforms other methods in
Table 2: Offline Reinforcement Learning Performance in MDP. 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 ( $\geq 0.95 \cdot \max$ ).
Dataset	Environment	BC	IQL	CQL	DT	TT	MOReL	DD	DD (reproduced)	Diffuser	DS (ours)
Medium-Expert	HalfCheetah	55.2	86.7	91.6	86.8	95.0	53.3	90.6	91.5 ±2.5	79.8	95.7 ±0.3
Medium-Expert	Hopper	52.5	91.5	105.4	107.6	110.0	108.7	111.8	111.6 ±2.8	107.2	107.0 ±3.2
Medium-Expert	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	42.6	47.4	44.0	42.6	46.9	42.1	49.1	46.4±5.1	44.2	47.8 ±0.4
Medium	Hopper	52.9	66.3	58.5	67.6	61.1	95.4	79.3	81.2 ±7.2	58.5	76.6 ±4.2
Medium	Walker2d	75.3	78.3	72.5	74.0	79.0	77.8	82.5	79.9±5.3	79.7	83.6 ±0.3
Medium-Replay	HalfCheetah	36.6	44.2	45.5	36.6	41.9	40.2	39.3	39.4±1.5	42.2	41.1 ±0.1
Medium-Replay	Hopper	18.1	94.7	95.0	82.7	91.5	93.6	100	95.3±3.7	96.8	89.5 ±4.2
Medium-Replay	Walker2d	26.0	73.9	77.2	66.6	82.6	49.8	75	72.3±3.1	61.2	80.7 ±1.5
Average	51.9	77.0	77.6	74.7	78.9	72.9	82.2	80.3	75.3	81.1
most environments and attains the highest average score of 74.3, which reflects a  $15.7\%$  performance improvement over the next best-performing approach Diffuser. This highlights the effectiveness of our approach in handling POMDP tasks by more expressively modeling the dependencies among observations, actions, and rewards.
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 ( $\geq 0.95 \cdot \max$ ).
Dataset	Environment	BC	DT	TT	DD	Diffuser	DS (ours)
Medium-Expert	HalfCheetah	42.1	80.8	94.9	19.07	82.2	92.7 ±0.8
Medium-Expert	Hopper	51.1	105.2	61.6	32.7	70.7	110.9 ±0.4
Medium-Expert	Walker2d	51.3	106.0	51.7	74.8	82.4	94.1 ±8.5
Medium	HalfCheetah	43.3	42.7	46.7	40.3	45.4	47.1 ±0.3
Medium	Hopper	36.4	63.1	55.7	38.1	62.2	57.7 ±3.9
Medium	Walker2d	39.4	64.2	28.5	53.2	55.7	74.3 ±4.2
Medium-Replay	HalfCheetah	2.1	35.5	43.8	39.8	39.3	40.3 ±1.2
Medium-Replay	Hopper	24.3	78.3	84.4	22.1	80.9	86.9 ±2.6
Medium-Replay	Walker2d	23.8	45.3	10.2	58.4	58.7	66.8 ±1.8
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-v2 POMDP using various reward and action models, with diffusion-based or transformer-based observation model. In each choice of observation model, the algorithm with the highest performance is highlighted.
Reward models	Action 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 ±1.7	56.9 ±2.2	36.3 ±3.1	58.0 ±4.4	46.9 ±3.7	34.9 ±3.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 of observations, rewards, and actions employing three distinct models that can be trained independently using teacher forcing. In this section, we explore the additional flexibility offered by different architecture choices for each module. For observation, reward, and action prediction, we consider diffusion models and Transformer-based autoregressive models. For reward and action models, we additionally consider MLPs that are restricted in their window and only look at the immediate state information to make a decision. The
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 suggest a clear advantage when incorporating reward modeling.
Observation model	Action models without reward modelling	Action models with reward modelling
Transformer	Diffusion	MLP	Transformer	Diffusion	MLP
Diffusion-based	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-based	45.1 ±5.2	39.4 ±3.2	39.6 ±3.7	58.0 ±4.4	54.3 ±3.3	42.5 ±4.1
results shown in Table 4 display a combination of  $2 \times 3 \times 3$  policy agents for the Hopper-medium v2 POMDP environment. Since we adopt a modular structure, we can compose the different modules efficiently and hence, we only needed to train 2 (state) + 3 (reward) + 3 (action) models. In Table 4, we find that the performance of pure transformer- or diffusion-based Decision Stacks gives reasonable performance (transformers: 53.0, diffusion: 56.9) but these pure combinations can be slightly outperformed by hybrids, e.g., the best achieving entry (58.2) in Table 4 uses a diffusion-based observation model, a transformer-based reward model and a diffusion-based action model. MLPs generally are outperformed by generative architectures, especially when used for modeling actions.
In addition to architectural flexibility, DS's modularity also supports compositional transfer across tasks. The lack of parameter sharing across modules not only enables hardware parallelization but also enables the transfer of modules to new environments. As shown in table 6, we created modified versions of the Maze2d-Umaze-v2 environment, altering action and reward spaces while keeping the observation model consistent. Both the reward and observation models were kept consistent during transfers between different action spaces. A single observation model was shared across six environments, and each of the two reward models was utilized across three action spaces. Our findings show that DS promotes efficient learning and compositional generalization, emphasizing its modular efficiency in scenarios with reusable components across varied tasks or environments.
Table 6: Experiments on Compositional Generalization. In the Maze2d-Umaze-v2 environment, we analyze both dense and sparse rewards. For the action space, we consider three variations of the force on a 2D ball: 1) Unscaled Force: the original force applied to the ball in the 2D maze environment; 2) Reflected Force: the original force mirrored across the x-axis; 3) Rotated Force: the original force rotated 30 degrees counterclockwise, altering the direction of the applied force. Each row shares the same reward model and observation model across the three different action spaces. The results underscore the modular efficiency of DS in scenarios where reusable components exist across different tasks or environments. The standard error is reported, and the results are averaged over 15 random seeds.
Action Space
Unscaled Force	Reflected Force	Rotated Force
Reward Space	dense	93.2 ± 10.7	90.3 ± 9.1	102.9 ± 8.1
sparse	111.3 ± 12.2	102.6 ± 3.9	115.2 ± 8.0
Furthermore, we compare the best reward modeling architectures with alternatives that do not consider rewards. This is standard practice for Diffuser [Janner et al., 2022] and Decision Diffuser (DD) [Ajay et al., 2022]. For example, DD predict the action at time  $t$  only based on the current observation and next observation,  $P(a_{t}|o_{t},o_{t + 1})$ , parameterized via an MLP. As delineated in Table 5, the inclusion of reward models significantly boosts performance in the dense reward POMDP environment. We include additional analysis and discussion in the Appendix D.
Offline Reinforcement Learning is a paradigm that for learning RL policies directly from previously logged interactions of a behavioral policy. The key challenge is that any surrogate models trained on an offline dataset do not generalize well outside the dataset. Various strategies have been proposed to mitigate the challenges due to distribution shifts by constraining the learned policy to be conservative
and closely aligned with the behavior policy. These include learning a value function that strictly serves as a lower bound for the true value function in CQL [Kumar et al., 2020], techniques focus on uncertainty estimation such as Kumar et al. [2019], and policy regularization methods [Wu et al., 2019, Ghasemipour et al., 2021, Kumar et al., 2019, Fujimoto and Gu, 2021, Fujimoto et al., 2019]. Model-based methods like MORel [Kidambi et al., 2020] and ReBeL [Lee et al., 2021] add pessimism into the dynamics models. In the context of partially observed settings, Rafailov et al. [2021] extends model-based offline RL algorithms by incorporating a latent-state dynamics model for high-dimensional visual observation spaces, effectively representing uncertainty in the latent space and Zheng et al. [2023] derive algorithms for offline RL in settings where trajectories might have missing actions. Our work takes the RL as inference perspective [Levine, 2018] and employs the tools of probabilistic reasoning and neural networks for training RL agents.
Generative models for offline RL. Over the past few years, the RL community has seen a growing interest in employing generative models for context-conditioned sequence generation by framing the decision-making problem as a generative sequence prediction problem. Here, we expand on our discussion from §3 with additional context and references. Decision Transformer [Chen et al., 2021] and Trajectory Transformer [Janner et al., 2021] concurrently proposed the use of autoregressive transformer-based models [Radford et al., 2018] for offline RL in model-based and model-free setups. Online decision transformer [Zheng et al., 2022] further finetunes the offline pretrained policies in online environments through a sequence-level exploration strategy. GATO [Reed et al., 2022, Lee et al., 2022] and PEDA [Zhu et al., 2023] scale these models to multi-task and multi-objective settings. MaskDP Liu et al. [2022] shows that other self-supervised objectives such as masking can also enable efficient offline RL, especially in goal-conditioned settings. These advancements have also been applied to other paradigms in sequential decision making such as black-box optimization and experimental design [Nguyen and Grover, 2022, Krishnamoorthy et al., 2023a,b]. Recent works have shown that changing the generative model from transformer to a diffuser with guidance can improve performance in certain environments and also permit planning for model-based extensions [Janner et al., 2021, Ajay et al., 2022, Wang et al., 2022, Chen et al., 2022]. Dai et al. [2023] considers an extension where the state model is a pretrained text2image model and actions are extracted from consecutive image frames. As discussed in §3, these works make specific design choices that do not guarantee the modularity, flexibility, and expressivity ensured by our framework.
We proposed Decision Stacks, a modular approach for learning goal-conditioned policies using offline datasets. Decision Stacks comprises of 3 modules tasked with the prediction of observations, rewards, and actions respectively. In doing so, we strive for the twin benefits of expressivity through autoregressive conditioning across the modules and flexibility in generative design within any individual module. We showed its empirical utility across a range of offline RL evaluations for both MDP and POMDP environments, as well as long-horizon planning problems. In all these settings, Decision Stacks matches or significantly outperforms competing approaches while also offering significant flexibility in the choice of generative architectures and training algorithms.
Limitations and Future Work. Our experiments are limited to state-based environments and extending Decision Stacks to image-based environments is a promising direction for future work especially in light of the gains we observed for POMDP environments. We are also interested in exploring the benefits of a modular design for pretraining and transfer of modules across similar environments and testing their generalization abilities. Finally, online finetuning of Decision Stacks using techniques similar to Zheng et al. [2022] is also an exciting direction of future work.
This research is supported by a Meta Research Award.
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 4975-4982, 2019.
J. Sohl-Dickstein, E. Weiss, N. Maheswaranathan, and S. Ganguli. Deep unsupervised learning using nonequilibrium thermodynamics. In International Conference on Machine Learning, pages 2256-2265. PMLR, 2015.
A. Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A. N. Gomez, L. Kaiser, and I. Polosukhin. Attention is all you need. Advances in 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.06193, 2022.
G. Williams, A. Aldrich, and E. Theodorou. Model predictive path integral control using covariance variable importance sampling. arXiv preprint arXiv:1509.01149, 2015.
R. J. Williams and D. Zipser. A learning algorithm for continually running fully recurrent neural networks. Neural computation, 1(2):270-280, 1989.
C. Wu, S. Yin, W. Qi, X. Wang, Z. Tang, and N. Duan. Visual chatgpt: Talking, drawing and editing with visual foundation models. arXiv preprint arXiv:2303.04671, 2023.
Y. Wu and K. He. Group normalization. In Proceedings of the European conference on computer vision (ECCV), pages 3-19, 2018.
Y. Wu, G. Tucker, and O. Nachum. Behavior regularized offline reinforcement learning. arXiv preprint arXiv:1911.11361, 2019.
Q. Zheng, A. Zhang, and A. Grover. Online decision transformer. In International Conference on 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 A. Grover. Scaling pareto-efficient decision making via offline multi-objective rl. In International Conference on Learning Representations, 2023.
A Model details of Decision Stacks
In Decision Stacks, we employ a modular approach by utilizing separate models for predicting observations, rewards, and actions. These models are trained independently using the technique of teacher forcing [Williams and Zipser, 1989]. This modular design allows us to explore the architectural flexibility of Decision Stacks and incorporate various inductive biases tailored to each modality. Below, we provide a description of the models used for each module in our approach:
A.1 Observations models
- Diffusion-based model. We adopt the diffusion process as a conditional generative model to generate future observations based on current and past observations. Following Decision Diffuser [Ajay et al., 2022], we employ classifier-free guidance [Ho and Salimans, 2022] and low-temperature sampling to generate future observations conditioned on goals. With classifier-free guidance, we first sample  $\boldsymbol{x}_K(\tau)$  from a Gaussian noise and refine it to  $\boldsymbol{x}_0(\tau)$  with the following procedure in Eq 7 for denoising  $\boldsymbol{x}_k(\tau)$  into  $\boldsymbol{x}_{k - 1}(\tau)$ .
$$
\hat {\epsilon} := \epsilon_ {\theta} (\boldsymbol {x} _ {k} (\tau), \varnothing , k) + \omega \left(\epsilon_ {\theta} (\boldsymbol {x} _ {k} (\tau), G, k) - \epsilon_ {\theta} (\boldsymbol {x} _ {k} (\tau), \varnothing , k)\right) \tag {7}
$$
where  $G$  is the goal information. For the observation model, we diffuse over a sequence of consecutive observations.
$$
\tau_ {o b s} = \left[ \begin{array}{l l l l} o _ {0} & o _ {1} & \dots & o _ {T} \end{array} \right] \tag {8}
$$
We parameterize the diffusion model as U-Net [Ajay et al., 2022, Janner et al., 2022]. To generate future observations, we use an inpainting strategy where we condition the first observation in the diffusion sequence as the current observation  $o_0$ .
- Transformer-based model. The trajectory sequences can also be autoregressively predicted using a transformer architecture as previously demonstrated in Decision Transformer [Chen et al., 2021] and Trajectory Transformer [Janner et al., 2021]. In our approach, we employ a GPT [Radford et al., 2018] model as the underlying transformer architecture. This transformer-based model is trained by feeding it with a sequence of observations, alongside the task goal as the conditioning information. To ensure temporal consistency during both training and testing, we employ causal self-attention masks, enabling the model to predict the next observation solely based on the preceding observations and task information. To effectively represent the observation and task modalities, we treat them as separate entities and embed them into dedicated tokens. These tokens are then input to the transformer, which is further enriched with positional encoding to capture the temporal relationships among the observations. This combination of a GPT model, self-attention masks, and modality-specific token embeddings allows our approach to effectively model the sequential nature of the observations while incorporating task-related information. The transformer can autoregressively generate a sequence of observations as in Eq 8.
A.2 Reward models
- Diffusion-based model. For our reward model, we employ a similar U-Net architecture as the observation diffusion model but diffuse over different sequences. Specifically, we diffuse over the combined sequence of observations and rewards as follows:
$$
\tau_ {r e w} = \left[ \begin{array}{c c c c} o _ {0} & o _ {1} & \dots & o _ {T} \\ r _ {0} & r _ {1} & \dots & r _ {T} \end{array} \right] \tag {9}
$$
In this combined sequence, we concatenate the observation sequence  $o_0, o_1, \dots, o_T$  and the corresponding reward sequence  $r_0, r_1, \dots, r_T$  at each time step for training. To generate the reward sequence, we utilize an inpainting conditioning strategy, similar to the one used in the observation model. This strategy involves conditioning the diffusion process on the observation sequences  $o_0, o_1, \dots, o_T$  while generating the rewards. By incorporating this inpainting conditioning, the reward model can effectively utilize the available observation information to generate accurate reward predictions throughout the diffusion process.