SlowGuess/NeurIPS_2023_OCR huggingface dataset:
- Transformer-based model. We employ an Encoder-Decoder transformer architecture for reward sequence generation, given a sequence of observations and the task goal. The encoder module embeds the observations and task goal, incorporating time encoding for capturing temporal dependencies. The encoded inputs are then processed by a transformer layer. The decoder module generates the reward sequence by iteratively predicting the next reward based on the encoded inputs and previously generated rewards. The transformer architecture facilitates capturing long-range dependencies and effectively modeling the dynamics between observations and rewards.
- MLP reward model. The multi-layer perceptron (MLP) reward model is a straightforward mapping from the current observation to the corresponding immediate reward. This model does not incorporate context or future synthesized information as it relies on a fixed input and output size. Consequently, the MLP architecture does not consider or incorporate any contextual information during its prediction process.
A.3 Action models
- Diffusion-based model. Similar to the observation and reward diffusion processes, we perform diffusion over the trajectory defined in Equation 10 as well as goal conditioning. By employing diffusion in this manner, we generate the action sequence conditioned on the information contained within the previous observations, rewards, and anticipated future observations and rewards. This diffusion process enables us to effectively capture the dependencies and dynamics among these elements, resulting in the generation of contextually informed action sequences.
$$
\tau_ {a c t} = \left[ \begin{array}{l l l l} o _ {0} & o _ {1} & \dots & o _ {T} \\ r _ {0} & r _ {1} & \dots & r _ {T} \\ a _ {0} & a _ {1} & \dots & a _ {T} \end{array} \right] \tag {10}
$$
- Transformer-based model. We employ the Encoder-Decoder transformer architecture, similar to that utilized in the reward transformer model, for our action model. The encoder module performs embedding of the observations, rewards, and task goals, incorporating time encodings to capture temporal dependencies. The encoded inputs are subsequently processed by a transformer layer. The decoder module is responsible for generating the action sequence by iteratively predicting the subsequent action based on the encoded inputs and previously generated actions. Consequently, our action model produces coherent and contextually informed action sequences.
- MLP action model. The multi-layer perceptron (MLP) action model functions as a direct mapping from the current observation and subsequent observation to the corresponding immediate action, similar to the inverse dynamics model in Decision Diffuser [Ajay et al., 2022]. However, due to its fixed input and output size, this MLP architecture does not incorporate contextual or future synthesized information. Consequently, the model lacks the ability to consider or integrate contextual details. This limitation proves disadvantageous in the context of partially observable Markov decision processes (POMDPs), where the inclusion of contextual information is vital for inferring hidden states and making informed decisions.
B Hyperparameters and training details
- The diffusion-based observation model follows the same training hyperparameters as in Decision Diffuser [Ajay et al., 2022], where we set the number of diffusion steps as 200 and the planning horizon as 100.
- For other models, we use a batch size of 32, a learning rate of  $3e^{-4}$ , and training steps of  $2e^{6}$  with Adam optimizer [Kingma and Ba, 2015].
- The MLP action model and the MLP reward model is a two layered MLP with 512 hidden units and ReLU activations.
- The diffusion models' noise model backbone is a U-Net with six repeated residual blocks. Each block consists of two temporal convolutions, each followed by group norm [Wu and He, 2018], and a final Mish nonlinearity [Misra, 2019].
- For Maze2D experiments, different mazes require different average episode steps to reach to target, we use the planning horizon of 180 for umaze, 256 for medium-maze and 300 for large maze.
- For Maze2D experiments, we use the warm-starting strategy where we perform a reduced number of forward diffusion steps using a previously generated plan as in Diffuser [Janner et al., 2022] to speed up the computation.
- The training of all three models, including the observation, action, and reward models, is conducted using the teacher-forcing technique [Williams and Zipser, 1989].
- Additional hyperparameters can be found in the configuration files within our codebase.
- Upon reproducing the Decision Diffuser (DD) approach [Ajay et al., 2022] using their provided codebase, we observed that the agent performance can be sensitive to the test return of the locomotion tasks, conditional guidance parameter, and sampling noise. For Decision Stacks variants that use the DD observation model, we directly use the models tuned for DD's best performance. In future work, it can be helpful to use conservative regularizers [Nguyen et al., 2022] to further improve both DD and DS performance.
Each model was trained on a single NVIDIA A5000 GPU.
C Sensitivity analysis on the dimension occlusion for POMDPs.
Table 7: 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 hopper-medium-expert-v2 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 2 dim of velocity	43.3 ± 2.2	59.0 ± 2.3	68.7 ± 8.3	52.3 ± 5.5	52.5 ± 7.4	68.0 ± 3.4
Occlude middle 2 dim of velocity	34.0 ± 0.9	73.1 ± 4.6	76.7 ± 7.9	31.6 ± 5.14	24.2 ± 3.1	81.1 ± 6.9
Occlude last 2 dim of velocity	51.1 ± 1.3	105.2 ± 1.0	61.6 ± 3.3	32.7 ± 2.4	70.7 ± 4.7	110.9 ± 0.4
Occlude first 2 dim of position	14.6 ± 2.4	40.4 ± 3.2	15.0 ± 4.1	8.7 ± 0.2	29.5 ± 1.7	38.9 ± 6.2
Average	35.7 ± 1.7	69.4 ± 2.8	55.5 ± 5.9	31.3 ± 3.3	44.2 ± 4.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 trained two sets of transformer-based action models: one set modeling the sequence of observations and actions, and the other set incorporating the sequence of observations, rewards, and actions. The evaluation of these models was conducted on both Markov Decision Process (MDP) and Partially Observable Markov Decision Process (POMDP) environments, specifically three locomotion tasks. The results obtained from this study highlight the evident advantages of reward modeling, particularly in dense-reward locomotion tasks. The details and findings are presented in the Table 8, emphasizing the benefits gained from incorporating reward modeling.
Table 8: Ablation on reward modeling.
Environments	Action models
Transformer with reward modelling	Transformer without reward modelling
MDP	Hopper-medium-v2	76.6 ±4.2	69.7 ±3.4
Walker2d-medium-v2	83.6 ±0.3	82.7 ±1.2
Halfcheetah-medium-v2	47.8 ±0.4	43.0 ±2.8
POMDP	Average (MDP envs)	69.3	65.1
Hopper-medium-v2	56.0 ±3.5	43.6 ±1.3
Walker2d-medium-v2	74.3 ±4.2	54.2 ±7.3
Halfcheetah-medium-v2	47.1 ±0.3	46.5 ±0.7
Average (POMDP envs)	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, we choose the ordering of observations, rewards, and actions. Specifically, we ordered observations prior to rewards to be consistent with the functional definitions in MDPs, where the reward is typically a function of observations (and potentially other variables), but not vice versa. We tested this choice empirically as well in early experiments, and found our choice to significantly outperform the counterpart. The table provides a comparison of two different orderings: Reward-State-Action (R, S, A) and State-Reward-Action (S, R, A). The results indicate that the S, R, A ordering outperforms the R, S, A ordering in the halfcheetah-medium-replay-v2 environment. DS models observations prior to rewards to be consistent with the functional definitions in MDPs, where the reward is typically a function of observations (and potentially other variables), but not vice versa.
Ordering	halfcheetah-medium-replay-v2 Performance
R, S, A	32.1 ± 0.1
S, R, A	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-based approach to generate future plans. In this task with goal-conditioning, the position of the maze's goal serves as the condition.
Regarding Diffuser, we strictly adhere to their codebase and observed that Diffuser impaints the last goal position into the diffusion sequence. Similarly, for DD, we follow their codebase and incorporate the goal as conditioning, embedding it and transmitting it to the diffusion backbone, which is an Unet. In the case of DD, we also apply inpainting conditioning by conditioning the last position in the diffusion sequence to serve as the goal position.
F.1 Detour problem of inpainting conditioning
However, we encountered an issue with this inpainting goal conditioning when performing open-loop generation. The problem arises because, during each replanning iteration of the diffusion model, it only conditions on the last time step to be the goal state. Consequently, when the agent is in close proximity to the goal point, the diffusion model plans a detour that initially takes the agent away from the goal and then redirects it back towards the goal. Since the environment's reward function is designed in such a way that the agent receives a reward when it is near the goal, and the episode does
not terminate upon reaching the goal point but rather when it reaches the maximum episode length, it is actually more advantageous for the agent to remain at the goal point.
After recognizing this issue, we devised a "progressive conditioning" (PC) strategy to address it. This approach involves gradually increasing the number of timesteps in the diffusion sequence that are conditioned to be the goal position as time progresses. By implementing this progressive conditioning method, we successfully resolve the detour problem and provide additional incentives for the agents to remain at the goal position. Without PC, the agent's behavior during open-loop evaluation exhibits a recurring pattern of moving toward the goal and deviating from it. This results in unfavorable trajectories where the agent fails to stay at the goal position.
The integration of progressive conditioning led to performance enhancement in both Diffuser and DD, although the overall results still indicate that Decision Stacks outperforms DD in both single goal (average performance: 109.8 vs. 131.5) and multi goal environments (average performance: 111.6 vs. 123.4). We compare the performance of Diffuser, DD, DS with and without progressive conditioning in the bar charts below in Figure 4 and Figure 5. The results indicate that DS with PC outperforms other baselines on most of the environment settings across single-goal and multi-goal settings.
Figure 4: Bar chart comparison of Diffuser, DD and DS with and without progressive conditioning in the single goal setting of Maze2D.
Figure 5: Bar chart comparison of Diffuser, DD and DS with and without progressive conditioning in the multi goal setting of Maze2D.