Original: D EEP T RANSFORMER Q-N ETWORKS FOR P ARTIALLY O BSERVABLE R EINFORCEMENT L EARNING
Revised: Deep Transformer Q-Networks for Partially Observable Reinforcement Learning

Original: Anonymous authors

Original: Paper under double-blind review

Original: A BSTRACT
Revised: Abstract

Original: Real-world reinforcement learning tasks often involve some form of partial observability where the observations only give a partial or noisy view of the true state of the world.
Revised: Real-world reinforcement learning tasks often involve some form of partial observability where the observations only give a partial or noisy view of the true state of the world.

Original: Such tasks typically require some form of memory, where the agent has access to multiple past observations, in order to perform well.
Revised: Such tasks typically require some form of memory, where the agent has access to multiple past observations, in order to perform well.

Original: One popular way to incorporate memory is by using a recurrent neural network to access the agent’s history.
Revised: One popular way to incorporate memory is by using a recurrent neural network to access the agent’s history.

Original: However, recurrent neural networks in reinforcement learning are often fragile and difficult to train and sometimes fail completely as a result.
Revised: However, recurrent neural networks in reinforcement learning are often fragile and difficult to train and sometimes fail completely as a result.

Original: In this work, we propose Deep Transformer Q-Networks (DTQN), a novel architecture utilizing transformers and self-attention to encode an agent’s history.
Revised: In this work, we propose Deep Transformer Q-Networks (DTQN), a novel architecture utilizing transformers and self-attention to encode an agent’s history.

Original: DTQN is designed modularly, and we compare results against several modifications to our base model.
Revised: DTQN is designed modularly, and we compare results against several modifications to our base model.

Original: Our experiments demonstrate that our approach can solve partially observable tasks faster and more stably than previous recurrent approaches.
Revised: Our experiments demonstrate that our approach can solve partially observable tasks faster and more stably than previous recurrent approaches.

Original:  I NTRODUCTION
Revised:  Introduction

Original: In recent years, deep neural networks have become the computational backbone of reinforcement learning, achieving strong performance across a wide array of difficult tasks including games (Mnih et al., 2015; Silver et al., 2016) and robotics (Levine et al., 2018; Gao et al., 2020).
Revised: In recent years, deep neural networks have become the computational backbone of reinforcement learning, achieving strong performance across a wide array of difficult tasks including games (Mnih et al., 2015; Silver et al., 2016) and robotics (Levine et al., 2018; Gao et al., 2020).

Original: In particular, Deep
Revised: In particular, Deep

Original: Q-Networks (DQN) (Mnih et al., 2015) revolutionized the field of deep RL by achieving super-human performance on Atari 2600 games in the Atari Learning Environment (Bellemare et al., 2013).
Revised: Q-Networks (DQN) (Mnih et al., 2015) revolutionized the field of deep RL by achieving super-human performance on Atari 2600 games in the Atari Learning Environment (Bellemare et al., 2013).

Original: Since then, several advancements have been proposed to improve DQN (Hessel et al., 2018), and deep RL has been shown to excel in continuous control tasks as well (Haarnoja et al., 2018; Fujimoto et al., 2018).
Revised: Since then, several advancements have been proposed to improve DQN (Hessel et al., 2018) and deep RL has been shown to excel in continuous control tasks as well (Haarnoja et al., 2018; Fujimoto et al., 2018).

Original: However, most Deep RL methods assume the agent is operating within a fully observable environment; that is, one in which the agent has access to the environment’s full state information.
Revised: However, most Deep RL methods assume the agent is operating within a fully observable environment; that is, one in which the agent has access to the environment’s full state information.

Original: But this assumption does not hold for many realistic domains due to components such as noisy sensors, occluded images, or additional unknown agents.
Revised: But this assumption does not hold for many realistic domains due to components such as noisy sensors, occluded images, or additional unknown agents.

Original: These domains are partially observable, and pose a much bigger challenge for RL compared to the standard fully observable setting.
Revised: These domains are partially observable, and pose a much bigger challenge for RL compared to the standard fully observable setting.

Original: Indeed, naïve methods often fail to learn in partially observable environments without additional architectural or training support (Pinto et al., 2017;
Revised: Indeed, naïve methods often fail to learn in partially observable environments without additional architectural or training support (Pinto et al., 2017;

Original: Igl et al., 2018; Ma et al., 2020).
Revised: Igl et al., 2018; Ma et al., 2020).

Original: To solve partially observable domains, RL agents may need to remember (some or possibly all) previous observations (Kaelbling et al., 1998).
Revised: To solve partially observable domains, RL agents may need to remember (some or possibly all) previous observations (Kaelbling et al., 1998).

Original: As a result, RL methods typically add some sort of memory component, allowing them to store or refer back to recent observations in order to make more informed decisions.
Revised: As a result, RL methods typically add some sort of memory component, allowing them to store or refer back to recent observations in order to make more informed decisions.

Original: The current state-of-the-art approaches integrate recurrent neural networks, like LSTMs (Hochreiter & Schmidhuber, 1997) or GRUs (Cho et al., 2014), in conjunction with fully observable Deep RL architectures to process an agent’s history (Ni et al., 2021).
Revised: The current state-of-the-art approaches integrate recurrent neural networks, like LSTMs (Hochreiter & Schmidhuber, 1997) or GRUs (Cho et al., 2014), in conjunction with fully observable Deep RL architectures to process an agent’s history (Ni et al., 2021).

Original: But recurrent neural networks (RNNs) can be fragile and difficult to train, often requiring complicated “warm-up” strategies to initialize its hidden state at the start of each training batch (Lample & Chaplot, 2017).
Revised: But recurrent neural networks (RNNs) can be fragile and difficult to train, often requiring complicated “warm-up” strategies to initialize its hidden state at the start of each training batch (Lample & Chaplot, 2017).

Original: Conversely, the Transformer has been shown to model sequences much better than RNNs and is ubiquitous in natural language processing (NLP) (Devlin et al., 2018) and increasingly common in computer vision (Dosovitskiy et al., 2020).
Revised: Conversely, the Transformer has been shown to model sequences much better than RNNs and is ubiquitous in natural language processing (NLP) (Devlin et al., 2018) and increasingly common in computer vision (Dosovitskiy et al., 2020).

Original: Therefore, we propose Deep Transformer Q-Network (DTQN), a novel architecture using selfattention to solve partially observable RL domains.
Revised: Therefore, we propose Deep Transformer Q-Network (DTQN), a novel architecture using selfattention to solve partially observable RL domains.

Original: DTQN leverages a transformer decoder architecture with learned positional encodings to represent an agent’s history and accurately predict Q-values at each timestep.
Revised: DTQN leverages a transformer decoder architecture with learned positional encodings to represent an agent’s history and accurately predict Q-values at each timestep.

Original: Rather than a standard approach that trains on a single next step for a given history, we propose a training regime called intermediate Q-value prediction, which allows us to train DTQN on the Q-values generated for each timestep in the agent’s observation history and provide more robust learning.
Revised: Rather than a standard approach that trains on a single next step for a given history, we propose a training regime called intermediate Q-value prediction, which allows us to train DTQN on the Q-values generated for each timestep in the agent’s observation history and provide more robust learning.

Original: DTQN encodes an agent’s history more effectively than recurrent methods, which we show empirically across several challenging partially observable environments.
Revised: DTQN encodes an agent’s history more effectively than recurrent methods, which we show empirically across several challenging partially observable environments.

Original: We evaluate and analyze several architectural components, including: gated skip connections (Parisotto et al., 2020), positional encodings, identity map reordering (Parisotto et al., 2020), and intermediate value prediction (Al-Rfou et al., 2019).
Revised: We evaluate and analyze several architectural components, including: gated skip connections (Parisotto et al., 2020), positional encodings, identity map reordering (Parisotto et al., 2020), and intermediate value prediction (Al-Rfou et al., 2019).

Original: Our results provide strong evidence that our approach can successfully represent agents’ histories in partially observable domains.
Revised: Our results provide strong evidence that our approach can successfully represent agents’ histories in partially observable domains.

Original: We visualize attention weights showing
Revised: We visualize attention weights showing

Original: DTQN learns an understanding of the domains as it works to solve tasks.
Revised: DTQN learns an understanding of the domains as it works to solve tasks.

Original:  B ACKGROUND
Revised:  Background

Original: When an environment does not emit its full state to the agent, the problem can be modeled as a Partially
Revised: When an environment does not emit its full state to the agent, the problem can be modeled as a Partially

Original: Observable Markov Decision Process (POMDP) Kaelbling et al.
Revised: Observable Markov Decision Process (POMDP) Kaelbling et al.

Original: A POMDP is formally described as the 6-tuple ( S , A , T , R , Ω , O ) .
Revised: A POMDP is formally described as the 6-tuple ( S , A , T , R , Ω , O ) .

Original: S , A , and Ω represent the environment’s set of states, actions, and observations, respectively.
Revised: S , A , and Ω represent the environment’s set of states, actions, and observations, respectively.

Original: T is the state transition function T ( s, a, s ′ )
Revised: T is the state transition function T ( s, a, s ′ )

Original: =
Revised: =

Original: P ( s ′ | s, a ) , denoting the probability of transitioning from state s to state s ′ given action a .
Revised: P ( s ′ | s, a ) , denoting the probability of transitioning from state s to state s ′ given action a .

Original: R describes the reward function R : S ×
Revised: R describes the reward function R : S ×

Original: A → R ; that is, the resultant scalar reward emitted by the environment for an agent that was in some state s ∈ S and took some action a ∈ A .
Revised: A → R ; that is, the resultant scalar reward emitted by the environment for an agent that was in some state s ∈ S and took some action a ∈ A .

Original: And O is the observation function
Revised: And O is the observation function

Original: O ( s ′ , a, o )
Revised: O ( s ′ , a, o )

Original: =
Revised: =

Original: P ( o | s ′ , a ) , the probability of observing o when action a is taken resulting in state s ′ .
Revised: P ( o | s ′ , a ) , the probability of observing o when action a is taken resulting in state s ′ .

Original: At each time step, t , the agent is in the environment’s state s t ∈ S , takes action a t ∈ A , manipulates the environment’s state to some s t +1 ∈ S based on the transition probability T ( s t , a t , s t +1 ) and receives a reward, r t =
Revised: At each time step, t , the agent is in the environment’s state s t ∈ S , takes action a t ∈ A , manipulates the environment’s state to some s t +1 ∈ S based on the transition probability T ( s t , a t , s t +1 ) and receives a reward, r t =

Original: R ( s t , a t ) .
Revised: R ( s t , a t ) .

Original: The goal of the agent is to maximize E (cid:2) (cid:80) t γ t r t (cid:3) , its expected discounted return for some discount factor γ ∈ [0 , 1) (Sutton & Barto, 2018).
Revised: The goal of the agent is to maximize E (cid:2) (cid:80) t γ t r t (cid:3) , its expected discounted return for some discount factor γ ∈ [0 , 1) (Sutton & Barto, 2018).

Original: Because agents in POMDPs do not have access to the environment’s full state information, they must rely on the observations o t ∈ Ω which relate to the state via the observation function,
Revised: Because agents in POMDPs do not have access to the environment’s full state information, they must rely on the observations o t ∈ Ω which relate to the state via the observation function,

Original: O ( s t +1 , a t , o t ) =
Revised: O ( s t +1 , a t , o t ) =

Original: P ( o t | s t +1 , a t ) .
Revised: P ( o t | s t +1 , a t ) .

Original: In general, agents acting in partially observable space cannot simply use observations as a proxy for state, since several states may be aliased into the same observation.
Revised: In general, agents acting in partially observable space cannot simply use observations as a proxy for state, since several states may be aliased into the same observation.

Original: Instead, they often consider some form of their full history of information, h t = { ( o 0 , a 0 ) , ( o 1 , a 1 ) , ..., ( o t − 1 , a t − 1 ) } .
Revised: Instead, they often consider some form of their full history of information, h t = { ( o 0 , a 0 ) , ( o 1 , a 1 ) , ..., ( o t − 1 , a t − 1 ) } .

Original: Because the history grows indefinitely as the agent proceeds in a trajectory, various ways of encoding the history exist.
Revised: Because the history grows indefinitely as the agent proceeds in a trajectory, various ways of encoding the history exist.

Original: Previous work has truncated the history to make it a fixed length (Zhu et al., 2017) or used an agent’s belief, which represents the estimate of the current state (Kaelbling et al., 1998).
Revised: Previous work has truncated the history to make it a fixed length (Zhu et al., 2017) or used an agent’s belief, which represents the estimate of the current state (Kaelbling et al., 1998).

Original: Since the deep learning revolution, others have used forms of recurrency, such as LSTMs and GRUs, to encode the history (Hausknecht & Stone, 2015; Yang & Nguyen, 2021).
Revised: Since the deep learning revolution, others have used forms of recurrency, such as LSTMs and GRUs, to encode the history (Hausknecht & Stone, 2015; Yang & Nguyen, 2021).

Original: 2. D EEP RECURRENT Q- NETWORKS
Revised: 2.1 Deep recurrent Q-networks

Original: Q-Learning (Watkins & Dayan, 1992) aims to learn a function Q : S ×
Revised: Q-Learning (Watkins & Dayan, 1992) aims to learn a function Q : S ×

Original: A → R which represents the value of each state-action pair in an MDP.
Revised: A → R which represents the value of each state-action pair in an MDP.

Original: Given a state s , action a , reward r , next state s ′ , and learning rate α , the Q -function is updated with the equation
Revised: Given a state s , action a , reward r , next state s ′ , and learning rate α , the Q -function is updated with the equation

Original: [Equation]
Revised: [Equation]

Original: In more challenging domains, however, the state-action space of the environment is often too large to be able to learn an exact Q -value for each state-action pair.
Revised: In more challenging domains, however, the state-action space of the environment is often too large to be able to learn an exact Q -value for each state-action pair.

Original: Instead of learning a tabular Q-function,
Revised: Instead of learning a tabular Q-function,

Original: DQN (Mnih et al., 2015) learns an approximate Q -function featuring strong generalization capabilities over similar states and actions.
Revised: DQN (Mnih et al., 2015) learns an approximate Q -function featuring strong generalization capabilities over similar states and actions.

Original: DQN is trained to minimize the Mean Squared Bellman Error
Revised: DQN is trained to minimize the Mean Squared Bellman Error

Original: [Equation]
Revised: [Equation]

Original: where transition tuples of states, actions, rewards, and future states ( s, a, r, s ′ ) are sampled uniformly from a replay buffer, D , of past experiences while training.
Revised: where transition tuples of states, actions, rewards, and future states ( s, a, r, s ′ ) are sampled uniformly from a replay buffer, D , of past experiences while training.

Original: The target r + max a ′ ∈A Q ( s ′ , a ′ ; θ ′ ) invokes DQN’s target network (parameterized by θ ′ ), which lags behind the main network (parameterized by θ ) to produce more stable updates.
Revised: The target r + max a ′ ∈A Q ( s ′ , a ′ ; θ ′ ) invokes DQN’s target network (parameterized by θ ′ ), which lags behind the main network (parameterized by θ ) to produce more stable updates.

Original: However, in partially observable domains, DQN may not learn a good policy by simply replacing the network’s input from states to observations (i.e., an agent can often perform better by remembering some history).
Revised: However, in partially observable domains, DQN may not learn a good policy by simply replacing the network’s input from states to observations (i.e., an agent can often perform better by remembering some history).

Original: To address this challenge, Deep Recurrent Q-Networks (DRQN) (Hausknecht &
Revised: To address this challenge, Deep Recurrent Q-Networks (DRQN) (Hausknecht &

Original: Stone, 2015) incorporated histories into the Q -function by way of a long short-term memory (LSTM) layer (Hochreiter & Schmidhuber, 1997).
Revised: Stone, 2015) incorporated histories into the Q -function by way of a long short-term memory (LSTM) layer (Hochreiter & Schmidhuber, 1997).

Original: In DRQN’s training procedure, the sampled states are replaced with histories 
Revised: In DRQN’s training procedure, the sampled states are replaced with histories h t : t + k = { o t , o t +1 , ..., o t + k } from timestep t to step t + k , sampled randomly within each episode. The hidden state of the LSTM is zeroed at the start of each update.

Original: 2. T RANSFORMER DECODERS
Revised: 2.2 Transformer decoders

Original: The transformer architecture (Vaswani et al., 2017), originally introduced for natural language processing, stacks blocks of attention layers (Bahdanau et al., 2014) and is typically used to model sequential data.
Revised: The transformer architecture (Vaswani et al., 2017), originally introduced for natural language processing, stacks blocks of attention layers (Bahdanau et al., 2014) and is typically used to model sequential data.

Original: Intuitively, the transformer’s attention module receives as input a sequence of tokens (e.g., a sequence of observations in an episode) and the model learns to place stronger weights or more attention on the most important tokens.
Revised: Intuitively, the transformer’s attention module receives as input a sequence of tokens (e.g., a sequence of observations in an episode) and the model learns to place stronger weights or more attention on the most important tokens.

Original: For more details about the attention module in transformers, refer to Appendix A.3.
Revised: For more details about the attention module in transformers, refer to Appendix A.3.

Original: While the original transformer architecture formed an encoder-decoder structure, recent works often use either the encoder (Devlin et al., 2018) or the decoder (Radford et al., 2018).
Revised: While the original transformer architecture formed an encoder-decoder structure, recent works often use either the encoder (Devlin et al., 2018) or the decoder (Radford et al., 2018).

Original: The key difference between the two is that the decoder applies a causal masking to the attention layer; that is, the ith token cannot attend to tokens which come later in the sequence.
Revised: The key difference between the two is that the decoder applies a causal masking to the attention layer; that is, the ith token cannot attend to tokens which come later in the sequence.

Original: In general, the transformer decoder has been shown to perform better on generative tasks like next token prediction, while the transformer encoder is able to learn richer representations and excels on tasks such as language understanding.
Revised: In general, the transformer decoder has been shown to perform better on generative tasks like next token prediction, while the transformer encoder is able to learn richer representations and excels on tasks such as language understanding.

Original: DTQN utilizes the transformer decoder structure.
Revised: DTQN utilizes the transformer decoder structure.

Original: Given a tensor of shape ( B, C, D ) , where B is the batch size, C is the context length, and D is the model’s dimensionality size, the transformer decoder layer returns a tensor of the same shape, enabling us to stack layers on top of each other.
Revised: Given a tensor of shape ( B, C, D ) , where B is the batch size, C is the context length, and D is the model’s dimensionality size, the transformer decoder layer returns a tensor of the same shape, enabling us to stack layers on top of each other.

Original: The last transformer layer’s output can then be projected to the desired shape, or sent as input to another network.
Revised: The last transformer layer’s output can then be projected to the desired shape, or sent as input to another network.

Original: To ensure the raw inputs are of the correct shape, we often prepend a feature extraction step, such as a lookup embedding for text or integers, a multilayer perceptron for vectors, or convolutional neural network for images.
Revised: To ensure the raw inputs are of the correct shape, we often prepend a feature extraction step, such as a lookup embedding for text or integers, a multilayer perceptron for vectors, or convolutional neural network for images.

Original:  R ELATED WORK
Revised:  Related work

Original: Since its inception, several works have built upon DRQN.
Revised: Since its inception, several works have built upon DRQN.

Original: For example, DRQN was shown to beat human test subjects on the challenging 3D VizDoom video game environment (Kempka et al., 2016) when augmented with game feature information (Lample & Chaplot, 2017).
Revised: For example, DRQN was shown to beat human test subjects on the challenging 3D VizDoom video game environment (Kempka et al., 2016) when augmented with game feature information (Lample & Chaplot, 2017).

Original: Action-based DRQN (ADRQN)
Revised: Action-based DRQN (ADRQN)

Original: (Zhu et al., 2017) conditioned the network on the agent’s full history rather than just its observation history.
Revised: (Zhu et al., 2017) conditioned the network on the agent’s full history rather than just its observation history.

Original: Deep Distributed Recurrent Q-Networks (DDRQN)
Revised: Deep Distributed Recurrent Q-Networks (DDRQN)

Original: (Foerster et al., 2016) extends
Revised: (Foerster et al., 2016) extends

Original: DRQN into the multi-agent reinforcement learning setting.
Revised: DRQN into the multi-agent reinforcement learning setting.

Original: Like ADRQN, DDRQN conditioned on actions, but also shared weights between each agent, all while forgoing each agent’s replay buffer to sample from.
Revised: Like ADRQN, DDRQN conditioned on actions, but also shared weights between each agent, all while forgoing each agent’s replay buffer to sample from.

Original: The concept of attention is also well-studied in the reinforcement learning setting.
Revised: The concept of attention is also well-studied in the reinforcement learning setting.

Original: The most closely related work to ours using attention in deep Q-learning is Deep Attention Recurrent Q-Network (DARQN)
Revised: The most closely related work to ours using attention in deep Q-learning is Deep Attention Recurrent Q-Network (DARQN)

Original: (Sorokin et al., 2015), which used attention to aid an LSTM’s representation of the agent’s history.
Revised: (Sorokin et al., 2015), which used attention to aid an LSTM’s representation of the agent’s history.

Original: Similarly, visual attention has been added to DRQN-like architectures in an effort towards creating more interpretable reinforcement learning algorithms (Mott et al., 2019).
Revised: Similarly, visual attention has been added to DRQN-like architectures in an effort towards creating more interpretable reinforcement learning algorithms (Mott et al., 2019).

Original: Unlike our work, which uses self-attention such that agent’s history forms the queries, keys, and values, these works use the recurrent network’s last output state to form the queries, and the environment’s most recent observation forms the keys and values.
Revised: Unlike our work, which uses self-attention such that agent’s history forms the queries, keys, and values, these works use the recurrent network’s last output state to form the queries, and the environment’s most recent observation forms the keys and values.

Original: In the multi-agent setting, Multi-Actor-Attention-Critic (Iqbal & Sha, 2019) created an attention module in which each agent’s query is their own observation, and the keys and values are formed by the other agents’ observations.
Revised: In the multi-agent setting, Multi-Actor-Attention-Critic (Iqbal & Sha, 2019) created an attention module in which each agent’s query is their own observation, and the keys and values are formed by the other agents’ observations.

Original: Finally, the Simple Neural
Revised: Finally, the Simple Neural

Original: AttentIve Learner (SNAIL) (Mishra et al., 2017) used attention to develop a meta-learning agent capable of transferring its skills to similar but different environments.
Revised: AttentIve Learner (SNAIL) (Mishra et al., 2017) used attention to develop a meta-learning agent capable of transferring its skills to similar but different environments.

Original: The use of transformers in reinforcement learning has become more popular within the last few years.
Revised: The use of transformers in reinforcement learning has become more popular within the last few years.

Original: In the offline reinforcement learning setting, Decision Transformer (Chen et al., 2021) and
Revised: In the offline reinforcement learning setting, Decision Transformer (Chen et al., 2021) and

Original: Trajectory Transformer (Janner et al., 2021) concurrently proposed the idea of using transformer decoders for sequence modeling, surpassing the current offline RL state of the art.
Revised: Trajectory Transformer (Janner et al., 2021) concurrently proposed the idea of using transformer decoders for sequence modeling, surpassing the current offline RL state of the art.

Original: Online Decision
Revised: Online Decision

Original: Transformer (Zheng et al., 2022) extended Decision Transformer (Chen et al., 2021) by treating the offline training as a pre-training step, and fine-tuned the transformer in the online setting for even better performance.
Revised: Transformer (Zheng et al., 2022) extended Decision Transformer (Chen et al., 2021) by treating the offline training as a pre-training step, and fine-tuned the transformer in the online setting for even better performance.

Original: FlexiBiT (Carroll et al., 2022) trained a transformer encoder to learn a variety of inference tasks, such as behavioural cloning, forward and backward modeling, and inferring an agent’s history given its current state.
Revised: FlexiBiT (Carroll et al., 2022) trained a transformer encoder to learn a variety of inference tasks, such as behavioural cloning, forward and backward modeling, and inferring an agent’s history given its current state.

Original: Contrary to our work, which is trained completely online using reinforcement learning, these works are specialized to take advantage of an offline RL dataset, and train their agents in a supervised way (Schmidhuber, 2019).
Revised: Contrary to our work, which is trained completely online using reinforcement learning, these works are specialized to take advantage of an offline RL dataset, and train their agents in a supervised way (Schmidhuber, 2019).

Original: Other methods use transformers to learn from scratch in the online RL setting, like GTrXL (Parisotto et al., 2020), which modifies the ordering of components within the transformer block, and introduces a new gating mechanism to replace the residual skip connections.
Revised: Other methods use transformers to learn from scratch in the online RL setting, like GTrXL (Parisotto et al., 2020), which modifies the ordering of components within the transformer block, and introduces a new gating mechanism to replace the residual skip connections.

Original: We compare the effects of these modifications to our architecture.
Revised: We compare the effects of these modifications to our architecture.

Original: Lightweight transformers have shown strong performance in text adventure games (Xu et al., 2020), and the transformer encoder was applied to video games (Upadhyay et al., 2019).
Revised: Lightweight transformers have shown strong performance in text adventure games (Xu et al., 2020), and the transformer encoder was applied to video games (Upadhyay et al., 2019).

Original: In contrast to these works, we utilize a multi-layer transformer decoder architecture.
Revised: In contrast to these works, we utilize a multi-layer transformer decoder architecture.

Original: Vision transformers have been used in conjuction with DQN to stabilize Q-learning with data augmentation, replacing the standard convolutional neural networks (Hansen et al., 2021).
Revised: Vision transformers have been used in conjuction with DQN to stabilize Q-learning with data augmentation, replacing the standard convolutional neural networks (Hansen et al., 2021).

Original: The self-attention block in their work attends to features within a single observation whereas ours attends throughout the agent’s history.
Revised: The self-attention block in their work attends to features within a single observation whereas ours attends throughout the agent’s history.

Original:  D EEP TRANSFORMER Q- NETWORK ARCHITECTURE
Revised:  Deep transformer Q-network architecture

Original: Transformers seem like a natural fit to represent histories in POMDPs, but there are several open questions regarding how to use them best in deep RL.
Revised: Transformers seem like a natural fit to represent histories in POMDPs, but there are several open questions regarding how to use them best in deep RL.

Original: In particular, it is unclear what form of transformer to use, how to integrate it into deep RL methods and how they should be trained.
Revised: In particular, it is unclear what form of transformer to use, how to integrate it into deep RL methods and how they should be trained.

Original: We chose to build DTQN using a transformer decoder structure incorporating learned position encodings, and train on the Q-values generated for each timestep in the agent’s observation history.
Revised: We chose to build DTQN using a transformer decoder structure incorporating learned position encodings, and train on the Q-values generated for each timestep in the agent’s observation history.

Original: DTQN takes as input the agent’s previous k observations, h t : t + k
Revised: DTQN takes as input the agent’s previous k observations, h t : t + k

Original: = { o t , o t +1 , ..., o t + k − 1 } , linearly projects each observation into the dimensionality of the model, and adds positional encodings to add informationabout the absolute temporal location of each observation. The embedded history is then passed through N transformer layers, and finally projected to the action space of the environment (see Figure 1 and Algorithm 1). DTQN outputs a set of Q-values relating to each observation in the input.
Revised: = { o t , o t +1 , ..., o t + k − 1 } , linearly projects each observation into the dimensionality of the model, and adds positional encodings to add information

Original: Algorithm function F ORWARD P ASS ( h t : t + k = { o t , o t +1 , ..., o t + k − 1 } ) )
Revised: Algorithm function F ORWARD P ASS ( h t : t + k = { o t , o t +1 , ..., o t + k − 1 } ) )

Original: E 

Revised: E

Original: =
Revised: =

Original: Embedding ( h t : t + k ) + Pos for Layer L = 1 , ..., N do
Revised: Embedding ( h t : t + k ) + Pos for Layer L = 1 , ..., N do

Original: Q L − 1 =
Revised: Q L − 1 =

Original: E L − 1 W QL − 1 , K L − 1 =
Revised: E L − 1 W QL − 1 , K L − 1 =

Original: E L − 1 W KL − 1 , V L − 1 =
Revised: E L − 1 W KL − 1 , V L − 1 =

Original: E L − 1 W VL −
Revised: E L − 1 W VL −

Original: E L = LayerNorm L 1 ( Combine L 1 ( MultiHeadAttention L ( Q L − 1 , K L − 1 , V L − 1 ) , E L − 1 ))
Revised: E L = LayerNorm L 1 ( Combine L 1 ( MultiHeadAttention L ( Q L − 1 , K L − 1 , V L − 1 ) , E L − 1 ))

Original: E L ← LayerNorm L 2 ( Combine L 2 ( FFN L ( E L ) , E L )) end for
Revised: E L ← LayerNorm L 2 ( Combine L 2 ( FFN L ( E L ) , E L )) end for

Original: Output ← FFN N ( E N ) ▷ Project output to action space end function function T RAIN
Revised: Output ← FFN N ( E N ) ▷ Project output to action space end function function T RAIN

Original: Sample a minibatch of contexts ( h t : t + k , a t : t + k , r t : t + k , h t +1: t + k +1 ) from replay buffer D for i = 1 , ..., k do
Revised: Sample a minibatch of contexts ( h t : t + k , a t : t + k , r t : t + k , h t +1: t + k +1 ) from replay buffer D for i = 1 , ..., k do

Original: L i ( θ ) = E ( . ) ∼ D (cid:2)(cid:0)
Revised: L i ( θ ) = E ( . ) ∼ D (cid:2)(cid:0)

Original: r t + i − 1 + max a ′ ∈A Q ( h t +1:
Revised: r t + i − 1 + max a ′ ∈A Q ( h t +1:

Original: t + i +1 , a ′ ; θ ′ )
Revised: t + i +1 , a ′ ; θ ′ )

Original: − Q ( h t : t
Revised: − Q ( h t : t

Original: + i , a t + i − 1 ; θ ) (cid:1) 2 (cid:3) end for end function function U PDATE θ ←
Revised: + i , a t + i − 1 ; θ ) (cid:1) 2 (cid:3) end for end function function U PDATE θ ←

Original: [Equation]
Revised: θ − α ∇ θ (cid:80) k i =

Original:  ) end function 
Revised: 1 L i ( θ ) end function about the absolute temporal location of each observation. The embedded history is then passed through N transformer layers, and finally projected to the action space of the environment (see Figure 1 and Algorithm 1). DTQN outputs a set of Q-values relating to each observation in the input.

Original: While we only use the Q-values from the most recent observation during execution, we train the network using all generated Q-values, even those relating to the observations at the beginning of the subhistory using the loss function in Algorithm 1.
Revised: While we only use the Q-values from the most recent observation during execution, we train the network using all generated Q-values, even those relating to the observations at the beginning of the subhistory using the loss function in Algorithm 1.

Original: This training regime challenges the network to predict the Q-values in situations where it has little to no context, and produces a more robust agent.
Revised: This training regime challenges the network to predict the Q-values in situations where it has little to no context, and produces a more robust agent.

Original: The remainder of this section expands on each contribution of the DTQN architecture.
Revised: The remainder of this section expands on each contribution of the DTQN architecture.

Original: 4. O BSERVATION EMBEDDINGS AND POSITIONAL ENCODINGS
Revised: 4.1 Observation embeddings and positional encodings

Original: Before the observation history is passed to DTQN’s transformer layers, each observation in the agent’s most recent k observations, h t : t + k , is linearly projected to the dimensionality of the transformer via a learned observation embedding (see Figure 1).
Revised: Before the observation history is passed to DTQN’s transformer layers, each observation in the agent’s most recent k observations, h t : t + k , is linearly projected to the dimensionality of the transformer via a learned observation embedding (see Figure 1).

Original: After embedding, we add a learned positional encoding to each observation based on its position in the observation history.
Revised: After embedding, we add a learned positional encoding to each observation based on its position in the observation history.

Original: This result, which we call E 0 in Algorithm 1, is the input to the first transformer layer in DTQN.
Revised: This result, which we call E 0 in Algorithm 1, is the input to the first transformer layer in DTQN.

Original: Position encodings are common practice in transformers, especially for NLP tasks, where they are well studied (Wang & Chen, 2020).
Revised: Position encodings are common practice in transformers, especially for NLP tasks, where they are well studied (Wang & Chen, 2020).

Original: However, the importance of position is less clear in the reinforcement learning setting.
Revised: However, the importance of position is less clear in the reinforcement learning setting.

Original: In some control tasks, the temporal position of an observation may not have any effect on its importance or meaning to solve the task.
Revised: In some control tasks, the temporal position of an observation may not have any effect on its importance or meaning to solve the task.

Original: For instance, the importance of the priest observation in the classic HeavenHell domain (Bonet, 1998) is not dependent on when the observation occurs in the episode.
Revised: For instance, the importance of the priest observation in the classic HeavenHell domain (Bonet, 1998) is not dependent on when the observation occurs in the episode.

Original: On the other hand, domains with more dynamic state transitions may benefit greatly from the positional information.
Revised: On the other hand, domains with more dynamic state transitions may benefit greatly from the positional information.

Original: For this reason, we choose to learn our positional encodings as it gives the agent the most flexibility in terms of how it chooses to use them.
Revised: For this reason, we choose to learn our positional encodings as it gives the agent the most flexibility in terms of how it chooses to use them.

Original: We ablate this choice by comparing our learned positional encodings to sinusoidal positional encodings (used in the original transformer (Vaswani et al., 2017)) as well as not using any positional encodings in section 5.4.
Revised: We ablate this choice by comparing our learned positional encodings to sinusoidal positional encodings (used in the original transformer (Vaswani et al., 2017)) as well as not using any positional encodings in section 5.4.

Original: 4. T RANSFORMER DECODER STRUCTURE
Revised: 4.2 Transformer decoder structure

Original: Like the original GPT architecture (Radford et al., 2018), each transformer layer in DTQN features two submodules: masked multi-headed self-attention and a position-wise feedforward network.
Revised: Like the original GPT architecture (Radford et al., 2018), each transformer layer in DTQN features two submodules: masked multi-headed self-attention and a position-wise feedforward network.

Original: As described in Algorithm 1, first we project the output of the previous layer, E L − 1 to the queries, Q , keys, K , and values, V , through the weight matrices W Q , W K , and W V , respectively.
Revised: As described in Algorithm 1, first we project the output of the previous layer, E L − 1 to the queries, Q , keys, K , and values, V , through the weight matrices W Q , W K , and W V , respectively.

Original: After each submodule, that submodule’s input and output are combined (see the “Combine” step in Figure 1) followed by a LayerNorm (Ba et al., 2016).
Revised: After each submodule, that submodule’s input and output are combined (see the “Combine” step in Figure 1) followed by a LayerNorm (Ba et al., 2016).

Original: Finally, after the last transformer layer, we project thefinal embedding ( E N in Algorithm 1) to the action space of our environment to represent the Q-value for each action.
Revised: Finally, after the last transformer layer, we project the

Original: [Figure]
Revised: [Figure]

Original: DTQN excels both in terms of learning speed as well as final performance, clearly outperforming the baselines on nearly all domains. Refer to section 5.2 for discussion of results.
Revised: Refer to section final embedding ( E N in Algorithm 1) to the action space of our environment to represent the Q-value for each action.

Original: DTQN uses a residual skip connection (He et al., 2016) to combine the two streams, matching the original transformer, in favor of other choices of combination layers such as the GRU gating combination layer Parisotto et al.
Revised: DTQN uses a residual skip connection (He et al., 2016) to combine the two streams, matching the original transformer, in favor of other choices of combination layers such as the GRU gating combination layer Parisotto et al.

Original: Another contested decision is the position of LayerNorm with respect to each submodule; the original transformer (Vaswani et al., 2017) and original GPT (Radford et al., 2018) apply LayerNorm after the combine step whereas other works have moved the
Revised: Another contested decision is the position of LayerNorm with respect to each submodule; the original transformer (Vaswani et al., 2017) and original GPT (Radford et al., 2018) apply LayerNorm after the combine step whereas other works have moved the

Original: LayerNorm to immediately before the submodule (Radford et al., 2019; Parisotto et al., 2020; Xu et al., 2020).
Revised: LayerNorm to immediately before the submodule (Radford et al., 2019; Parisotto et al., 2020; Xu et al., 2020).

Original: DTQN applies the LayerNorm after the combine step, a choice we found to be simple while also demonstrating strong empirical performance.
Revised: DTQN applies the LayerNorm after the combine step, a choice we found to be simple while also demonstrating strong empirical performance.

Original: We ablate our choices of network with the aforementioned variants in section 5.3.
Revised: We ablate our choices of network with the aforementioned variants in section 5.3.

Original: 4. I NTERMEDIATE Q- VALUE PREDICTION
Revised: 4.3 Intermediate Q-value prediction

Original: DTQN outputs a set of Q-values for each timestep in the agent’s observation history.
Revised: DTQN outputs a set of Q-values for each timestep in the agent’s observation history.

Original: During evaluation, DTQN selects the action with the highest Q-value from the last timestep in its history.
Revised: During evaluation, DTQN selects the action with the highest Q-value from the last timestep in its history.

Original: It would therefore be straightforward to train DTQN using just the last timestep’s Q-values, since those have the most context to work with and are the most informed to select the optimal action.
Revised: It would therefore be straightforward to train DTQN using just the last timestep’s Q-values, since those have the most context to work with and are the most informed to select the optimal action.

Original: This regime, however, is very wasteful, as only a fraction of the generated Q-values actually get used for training.
Revised: This regime, however, is very wasteful, as only a fraction of the generated Q-values actually get used for training.

Original: Instead, we train DTQN using all generated Q-values.
Revised: Instead, we train DTQN using all generated Q-values.

Original: Originally used in the NLP setting where each position was tasked with predicting the next character and formed an auxiliary loss (Al-Rfou et al., 2019), we adapt this training regime to the reinforcement learning setting, as shown in Algorithm 1.
Revised: Originally used in the NLP setting where each position was tasked with predicting the next character and formed an auxiliary loss (Al-Rfou et al., 2019), we adapt this training regime to the reinforcement learning setting, as shown in Algorithm 1.

Original: Note that the for loop depicted in Algorithm 1 can be done in one forward pass of the network because of the causally-masked self-attention mechanism.
Revised: Note that the for loop depicted in Algorithm 1 can be done in one forward pass of the network because of the causally-masked self-attention mechanism.

Original: We ablate training based on all Q-values with training only on the last timestep’s Q-values in section 5.5, and show the performance gains in Figure 1.
Revised: We ablate training based on all Q-values with training only on the last timestep’s Q-values in section 5.5, and show the performance gains in Figure 1.

Original:  E XPERIMENTS
Revised:  Experiments

Original: Our experimental evaluation is designed to compare DTQN not only to previous Q-network baselines, but also to ablate our own method with other architectural choices.
Revised: Our experimental evaluation is designed to compare DTQN not only to previous Q-network baselines, but also to ablate our own method with other architectural choices.

Original: We evaluate these methods on a range of challenging domains featuring partial observations and requiring memory to solve them.
Revised: We evaluate these methods on a range of challenging domains featuring partial observations and requiring memory to solve them.

Original: We baseline DTQN against Deep Recurrent Q-Networks (DRQN) (Hausknecht & Stone, 2015) to show the transformer is a more effective history representation module than RNNs, Deep
Revised: We baseline DTQN against Deep Recurrent Q-Networks (DRQN) (Hausknecht & Stone, 2015) to show the transformer is a more effective history representation module than RNNs, Deep

Original: Q-Networks (DQN) (Mnih et al., 2015) to demonstrate the need for memory to solve the task consistently, and against an attention baseline (called “ATTN” in Figure 2) to show the benefits of our architectural choices.
Revised: Q-Networks (DQN) (Mnih et al., 2015) to demonstrate the need for memory to solve the task consistently, and against an attention baseline (called “ATTN” in Figure 2) to show the benefits of our architectural choices.

Original: ATTN, like the transformer, has observation and position embeddings, attention and feedforward network modules, but does not have LayerNorm or skip connections, and does not stack multiple blocks.
Revised: ATTN, like the transformer, has observation and position embeddings, attention and feedforward network modules, but does not have LayerNorm or skip connections, and does not stack multiple blocks.

Original: 5. D OMAINS
Revised: 5.1 Domains

Original: [Figure]
Revised: [Figure]

Original: We conduct our experiments on 4 different environment sets designed to challenge DTQN in different ways: classic POMDPs, gym-gridverse (GV) (Baisero & Katt, 2021), car flag (Nguyen, 2021), and memory cards.
Revised: We conduct our experiments on 4 different environment sets designed to challenge DTQN in different ways: classic POMDPs, gym-gridverse (GV) (Baisero & Katt, 2021), car flag (Nguyen, 2021), and memory cards.

Original: Hallway (Littman et al., 1995) and HeavenHell (Bonet, 1998) are classic navigation POMDPs requiring the agent to take and remember several information gathering steps before it can consistently achieve its goal.
Revised: Hallway (Littman et al., 1995) and HeavenHell (Bonet, 1998) are classic navigation POMDPs requiring the agent to take and remember several information gathering steps before it can consistently achieve its goal.

Original: Gym-Gridverse offers procedurally generated gridworlds containing difficult partially observable tasks.
Revised: Gym-Gridverse offers procedurally generated gridworlds containing difficult partially observable tasks.

Original: The agent’s field of view is restricted such that it can only see the cells in a 2 × 3 grid in front of it (see Figure 3), which introduces state aliasing and forces the agent to gather localizing information before it can successfully solve the task.
Revised: The agent’s field of view is restricted such that it can only see the cells in a 2 × 3 grid in front of it (see Figure 3), which introduces state aliasing and forces the agent to gather localizing information before it can successfully solve the task.

Original: We evaluate our agents in gridverse environments “Memory” and “Memory Four Rooms”, which require the agent to first find the colored information beacons, and then go to the flag whose color matches the beacon.
Revised: We evaluate our agents in gridverse environments “Memory” and “Memory Four Rooms”, which require the agent to first find the colored information beacons, and then go to the flag whose color matches the beacon.

Original: The colors of the flags and beacons are initialized randomly and, in Memory-Four-Rooms, the locations of the flag and beacons are also initialized randomly, increasing the environment’s difficulty.
Revised: The colors of the flags and beacons are initialized randomly and, in Memory-Four-Rooms, the locations of the flag and beacons are also initialized randomly, increasing the environment’s difficulty.

Original: Example screenshots of the gridverse domains are shown in Figure 3.
Revised: Example screenshots of the gridverse domains are shown in Figure 3.

Original: Car Flag features a car on a 1D line, where the car must first drive to an oracle flag to learn which direction the finish line is.
Revised: Car Flag features a car on a 1D line, where the car must first drive to an oracle flag to learn which direction the finish line is.

Original: Memory cards is a novel domain designed to test how much information an agent can memorize.
Revised: Memory cards is a novel domain designed to test how much information an agent can memorize.

Original: Based on the popular children’s memory card game,pairs of cards are hidden to the agent, with one card revealed at each timestep and the agent must guess the position of that card’s pair.
Revised: Based on the popular children’s memory card game,pairs of cards are hidden to the agent, with one card revealed at each timestep and the agent must guess the position of that card’s pair.

Original: We chose this set of domains to be representative of challenging partially observable problems.
Revised: We chose this set of domains to be representative of challenging partially observable problems.

Original: For more information regarding the domains, please see Appendix B.
Revised: For more information regarding the domains, please see Appendix B.

Original: 5. B ASELINE COMPARISON
Revised: 5.2 Baseline comparison

Original: Our evaluation against baselines is shown in Figure 2.
Revised: Our evaluation against baselines is shown in Figure 2.

Original: Each learning curve shows the success rate of the agent in the environment, where the line is the mean across 5 random seeds, and the shaded region represents standard error.
Revised: Each learning curve shows the success rate of the agent in the environment, where the line is the mean across 5 random seeds, and the shaded region represents standard error.

Original: We search across hyperparameters of interest, and select the best performing hyperparameter set, prioritizing consistency across domains.
Revised: We search across hyperparameters of interest, and select the best performing hyperparameter set, prioritizing consistency across domains.

Original: For specific hyperparameters and training details, refer to Appendix A.
Revised: For specific hyperparameters and training details, refer to Appendix A.

Original: The results in Figure 2 show DTQN outperforms the baselines in terms of learning speed and final performance on nearly all domains.
Revised: The results in Figure 2 show DTQN outperforms the baselines in terms of learning speed and final performance on nearly all domains.

Original: ATTN learns quickly, but rarely reaches optimal performance, and often becomes unstable.
Revised: ATTN learns quickly, but rarely reaches optimal performance, and often becomes unstable.

Original: DRQN, featuring an LSTM as its memory module, often performs well in our set of domains, but learns slower than DTQN, and is in general less stable.
Revised: DRQN, featuring an LSTM as its memory module, often performs well in our set of domains, but learns slower than DTQN, and is in general less stable.

Original: Sometimes, especially in the gridverse memory domains, DRQN’s performance plummets shortly after it begins to learn.
Revised: Sometimes, especially in the gridverse memory domains, DRQN’s performance plummets shortly after it begins to learn.

Original: It then struggles to regain its initial performance gains, taking as many as one million timesteps in the Gridverse memory 7x7 domain to improve its success rate.
Revised: It then struggles to regain its initial performance gains, taking as many as one million timesteps in the Gridverse memory 7x7 domain to improve its success rate.

Original: DQN, designed for MDPs and without any form of memory in its architecture, fails to achieve higher than 50% success rate on any of our domains, exemplifying the difficulty of the domains and the importance of using memory to solve them.
Revised: DQN, designed for MDPs and without any form of memory in its architecture, fails to achieve higher than 50% success rate on any of our domains, exemplifying the difficulty of the domains and the importance of using memory to solve them.

Original: These results highlight the effectiveness of DTQN in solving a range of POMDPs.
Revised: These results highlight the effectiveness of DTQN in solving a range of POMDPs.

Original: 5. GRU- GATES AND IDENTITY MAP REORDERING
Revised: 5.3 GRU-gates and identity map reordering

Original: In this section, we compare DTQN with different forms of the “Combine“ step (see Figure 1) as well as different positions of LayerNorm.
Revised: In this section, we compare DTQN with different forms of the “Combine“ step (see Figure 1) as well as different positions of LayerNorm.

Original: DTQN’s combine step is a residual skip connection, and the LayerNorm occurs after both the attention and the feedforward submodules.
Revised: DTQN’s combine step is a residual skip connection, and the LayerNorm occurs after both the attention and the feedforward submodules.

Original: In contrast,
Revised: In contrast,

Original: GTrXL (Parisotto et al., 2020) introduced GRU-like gating in the combine step, and identity map reordering, which moves the LayerNorms directly in front of the masked multi self attention and feedforward sections.
Revised: GTrXL (Parisotto et al., 2020) introduced GRU-like gating in the combine step, and identity map reordering, which moves the LayerNorms directly in front of the masked multi self attention and feedforward sections.

Original: We compare our DTQN with residual skip connection to a version of DTQN which uses GRU-like gating, a version of DTQN which uses identity map reordering, and a version which uses both GRU-like gating and identity map reordering.
Revised: We compare our DTQN with residual skip connection to a version of DTQN which uses GRU-like gating, a version of DTQN which uses identity map reordering, and a version which uses both GRU-like gating and identity map reordering.

Original: When both GRU gates and the identity map reordering is used, the architecture resembles GTrXL.
Revised: When both GRU gates and the identity map reordering is used, the architecture resembles GTrXL.

Original: However, we do not use the
Revised: However, we do not use the

Original: TransformerXL
Revised: TransformerXL

Original: (Dai et al., 2019) as in GTrXL, therefore we are not comparing to an exact replica.
Revised: (Dai et al., 2019) as in GTrXL, therefore we are not comparing to an exact replica.

Original: The results for this ablation are shown in Table 1.
Revised: The results for this ablation are shown in Table 1.

Original: DTQN performs competitively with the ablated versions.
Revised: DTQN performs competitively with the ablated versions.

Original: The variant with only identity map reordering performs significantly worse than the other versions, and the version with both identity map reordering and GRU-like gating performs worst on hallway.
Revised: The variant with only identity map reordering performs significantly worse than the other versions, and the version with both identity map reordering and GRU-like gating performs worst on hallway.

Original: Both DTQN and the GRU-like gating variant perform competitively on all three domains we tested.
Revised: Both DTQN and the GRU-like gating variant perform competitively on all three domains we tested.

Original: Although we do not use the TransformerXL in our experiments, we would expect to see the same relative performance across if we replaced our transformer decoder with the TransformerXL.
Revised: Although we do not use the TransformerXL in our experiments, we would expect to see the same relative performance across if we replaced our transformer decoder with the TransformerXL.

Original: A comparison of different transformer backbones, such as Big Bird (Zaheer et al., 2020), sparse transformers (Child et al., 2019), or the TransformerXL would be interesting future study.
Revised: A comparison of different transformer backbones, such as Big Bird (Zaheer et al., 2020), sparse transformers (Child et al., 2019), or the TransformerXL would be interesting future study.

Original: 5. P OSITIONAL ENCODINGS
Revised: 5.4 Positional encodings

Original: DTQN uses learned positional encodings to allow the network to adapt to different domains.
Revised: DTQN uses learned positional encodings to allow the network to adapt to different domains.

Original: Partially observable domains will exhibit a broad range of temporal sensitivity, and we want to provide DTQN the flexibility to learn encodings to match its domain.
Revised: Partially observable domains will exhibit a broad range of temporal sensitivity, and we want to provide DTQN the flexibility to learn encodings to match its domain.

Original: In this section, we compare the use of learned positional encodings with the sinusoidal encodings in the original transformer (Vaswani et al., 2017) as well as no positional encodings.
Revised: In this section, we compare the use of learned positional encodings with the sinusoidal encodings in the original transformer (Vaswani et al., 2017) as well as no positional encodings.

Original: The results for this comparison are shown in Table 1.
Revised: The results for this comparison are shown in Table 1.

Original: In the memory cards domain, the variant of DTQN without positional encodings performs significantly worse than both our learned encodings as well as the sinusoidal encodings.
Revised: In the memory cards domain, the variant of DTQN without positional encodings performs significantly worse than both our learned encodings as well as the sinusoidal encodings.

Original: However, in the gridverse memory task and hallway, the three styles of positional encodings perform comparably.
Revised: However, in the gridverse memory task and hallway, the three styles of positional encodings perform comparably.

Original: We analyze the resulting learned positional encodings from our trained DTQN agents across various domains in
Revised: We analyze the resulting learned positional encodings from our trained DTQN agents across various domains in

Original: Appendix E.
Revised: Appendix E.

Original: 5. I NTERMEDIATE Q- VALUE PREDICTION
Revised: 5.5 Intermediate Q-value prediction

Original: DTQN predicts and trains on the Q-values generated for each timestep in the agent’s observation history.
Revised: DTQN predicts and trains on the Q-values generated for each timestep in the agent’s observation history.

Original: During evaluation, however, we only consider the last timestep’s Q-values when determining which action to take.
Revised: During evaluation, however, we only consider the last timestep’s Q-values when determining which action to take.

Original: We could, therefore, train in the same way, only training with the last timesteps’
Revised: We could, therefore, train in the same way, only training with the last timesteps’

Original: Q-values.
Revised: Q-values.

Original: We compare these two training regimes, and the results for this are shown in Table 1.
Revised: We compare these two training regimes, and the results for this are shown in Table 1.

Original: Our results show the variant trained without intermediate Q-values suffers a significant performance decrease.
Revised: Our results show the variant trained without intermediate Q-values suffers a significant performance decrease.

Original: In the memory cards case, DTQN excels and solves the task with nearly 90% success rate, but the variant without intermediate Q-value prediction can barely solve the task 10% of the time.
Revised: In the memory cards case, DTQN excels and solves the task with nearly 90% success rate, but the variant without intermediate Q-value prediction can barely solve the task 10% of the time.

Original: By training on all generated Q-values, we produce a more robust and effective agent.
Revised: By training on all generated Q-values, we produce a more robust and effective agent.

Original: [Table]
Revised: [Table]

Original:  D ISCUSSION
Revised:  Discussion

Original: DTQN outperforms or is competitive with our baselines in terms of learning speed and final performance on all our domains.
Revised: DTQN outperforms or is competitive with our baselines in terms of learning speed and final performance on all our domains.

Original: One additional advantage of transformers is the ability to visualize self-attention weights as a form of interpreting the model.
Revised: One additional advantage of transformers is the ability to visualize self-attention weights as a form of interpreting the model.

Original: Intuitively, the selfattention mechanism allows the agent to prioritize observations in its history which provide it with the most information useful in solving its task.
Revised: Intuitively, the self-attention mechanism allows the agent to prioritize observations in its history which provide it with the most information useful in solving its task.

Original: The causal masking ensures the agent cannot attend to observations in its future, resembling how the agent will need to perform during execution.
Revised: The causal masking ensures the agent cannot attend to observations in its future, resembling how the agent will need to perform during execution.

Original: While the use of attention weights as a tool for explainability is still being studied (Jain & Wallace, 2019; Wiegreffe & Pinter, 2019), it does allow us to observe which observations the agent finds most valuable in its history.
Revised: While the use of attention weights as a tool for explainability is still being studied (Jain & Wallace, 2019; Wiegreffe & Pinter, 2019), it does allow us to observe which observations the agent finds most valuable in its history.

Original: In Figure 4, we visualize a trained DTQN agent’s attention weights from a trajectory in gridverse memory 7x7.
Revised: In Figure 4, we visualize a trained DTQN agent’s attention weights from a trajectory in gridverse memory 7x7.

Original: Crucially, the observation including the green beacon (circle with X, magnified on right) is strongly attended to by all future observations, indicating the DTQN agent has correctly learned which observations are important in solving the task.
Revised: Crucially, the observation including the green beacon (circle with X, magnified on right) is strongly attended to by all future observations, indicating the DTQN agent has correctly learned which observations are important in solving the task.

Original: When the agent sees the green flag (magnified on left), it attends to the observation with green beacon to ensure it selects the correct flag.
Revised: When the agent sees the green flag (magnified on left), it attends to the observation with green beacon to ensure it selects the correct flag.

Original: We provide additional attention visualizations in Appendix D
Revised: We provide additional attention visualizations in Appendix D

Original:  C ONCLUSION
Revised:  Conclusion

Original: In this work we introduce Deep Transformer Q-Networks, a novel architecture for solving challenging partially observable domains with reinforcement learning.
Revised: In this work we introduce Deep Transformer Q-Networks, a novel architecture for solving challenging partially observable domains with reinforcement learning.

Original: DTQN incorporates the transformer decoder, which excels in generating Q-values for each timestep of the agent’s observation history.
Revised: DTQN incorporates the transformer decoder, which excels in generating Q-values for each timestep of the agent’s observation history.

Original: We train the model on all generated Q-values, enabling an efficient training regime and faster learning.
Revised: We train the model on all generated Q-values, enabling an efficient training regime and faster learning.

Original: DTQN also utilizes learned positional encodings, empowering the model to learn domain-specific encodings which match the temporal dependencies of the environment.
Revised: DTQN also utilizes learned positional encodings, empowering the model to learn domain-specific encodings which match the temporal dependencies of the environment.

Original: We explore and ablate several architectural structures, and find our choices either outperform or are at least competitive with all tested variants.
Revised: We explore and ablate several architectural structures, and find our choices either outperform or are at least competitive with all tested variants.

Original: Finally, we provide a modular code  implementation of DTQN that is easy to extend and modify, which we hope the research community will be able to use as we expect our approach to serve as the basis and benchmark for future transformer-based methods in partially observable reinforcement learning.
Revised: Finally, we provide a modular code 1 implementation of DTQN that is easy to extend and modify, which we hope the research community will be able to use as we expect our approach to serve as the basis and benchmark for future transformer-based methods in partially observable reinforcement learning.

Original: A I MPLEMENTATION DETAILS AND HYPERPARAMETERS A. I MPLEMENTATION
Revised: A Implementation details and hyperparameters A.1 Implementation

Original: We train DTQN off-policy.
Revised: We train DTQN off-policy.

Original: Before the agent begins training, the replay buffer is seeded with 50, timesteps generated by a random policy.
Revised: Before the agent begins training, the replay buffer is seeded with 50, timesteps generated by a random policy.

Original: At each timestep of training, DTQN receives the agent’s previous k observations, and outputs k sets of Q-values, one for each observation.
Revised: At each timestep of training, DTQN receives the agent’s previous k observations, and outputs k sets of Q-values, one for each observation.

Original: During training, the beahvior policy acts in an ϵ -greedy way; that is, with probability ϵ , the agent selects an action randomly, and with probability (1 − ϵ ) , selects the action with the highest predicted Q value.
Revised: During training, the beahvior policy acts in an ϵ -greedy way; that is, with probability ϵ , the agent selects an action randomly, and with probability (1 − ϵ ) , selects the action with the highest predicted Q value.

Original: We anneal ϵ linearly from 1 . 0 to 0 .
Revised: We anneal ϵ linearly from 1 . 0 to 0 .

Original: 1 throughout the first 10% of timesteps, and then hold it fixed at 0 .
Revised: 1 throughout the first 10% of timesteps, and then hold it fixed at 0 .

Original: 1 for the rest of training.
Revised: 1 for the rest of training.

Original: The target values used for training are selected via the double DQN ((Van Hasselt et al., 2016)) strategy.
Revised: The target values used for training are selected via the double DQN ((Van Hasselt et al., 2016)) strategy.

Original: Every 10,000 training timesteps, the target network updates its parameters by copying the policy network’s current parameters.
Revised: Every 10,000 training timesteps, the target network updates its parameters by copying the policy network’s current parameters.

Original: The evaluation policy behaves greedily and always selects the action with the highest predicted Q-value.
Revised: The evaluation policy behaves greedily and always selects the action with the highest predicted Q-value.

Original: We trained DTQN on a NVIDIA GeForce RTX 2070 GPU, where each run took about four hours per one million timesteps.
Revised: We trained DTQN on a NVIDIA GeForce RTX 2070 GPU, where each run took about four hours per one million timesteps.

Original: A. H YPERPARAMETERS
Revised: A.2 Hyperparameters

Original: The full list of hyperparameters is listed in Table 2.
Revised: The full list of hyperparameters is listed in Table 2.

Original: We prioritized consistency across domains, although in the Hallway , HeavenHell , and CarFlag domains, we set d model to 64 rather than 128 .
Revised: We prioritized consistency across domains, although in the Hallway , HeavenHell , and CarFlag domains, we set d model to 64 rather than 128 .

Original: [Table]
Revised: [Table]

Original: A. A TTENTION B ACKGROUND
Revised: A.3 Attention Background

Original: DTQN uses a variant of attention called multi-headed scaled dot-product self-attention.
Revised: DTQN uses a variant of attention called multi-headed scaled dot-product self-attention.

Original: Given a sequence of observations in an episode, { o i } , we first project the observations into the model’s dimensionality using an embedding layer.
Revised: Given a sequence of observations in an episode, { o i } , we first project the observations into the model’s dimensionality using an embedding layer.

Original: The attention mechanism then performs linear transformations against learnable weight matrices W Q , W K , and W V to map the sequence of tokens to a sequence of queries, Q , a set of keys, K , and a set of values, V .
Revised: The attention mechanism then performs linear transformations against learnable weight matrices W Q , W K , and W V to map the sequence of tokens to a sequence of queries, Q , a set of keys, K , and a set of values, V .

Original: Once we have the keys, queries, and values, we compute the attention as follows:
Revised: Once we have the keys, queries, and values, we compute the attention as follows:

Original: [Equation]
Revised: [Equation]

Original: where d k is the dimensionality of K .
Revised: where d k is the dimensionality of K .

Original: The result of the softmax above gives an attention score for each query.
Revised: The result of the softmax above gives an attention score for each query.

Original: By splitting the weight matrices W Q , W K , and W V into smaller components, we form several independent heads, hence the name multi-headed attention.
Revised: By splitting the weight matrices W Q , W K , and W V into smaller components, we form several independent heads, hence the name multi-headed attention.

Original: The intuition behind multi-headed attention is to give each head the ability to attend to different parts of the input.
Revised: The intuition behind multi-headed attention is to give each head the ability to attend to different parts of the input.

Original: The result of each [Table] result
Revised: The result of each head is concatenated before the final linear projection result

Original: MultiHead ( Q, K, V ) =
Revised: MultiHead ( Q, K, V ) =

Original: Concat ( head 1 , ..., head h ) W O where head i
Revised: Concat ( head 1 , ..., head h ) W O where head i

Original: =
Revised: =

Original: Attention ( QW Qi , KW Ki , V W Vi )
Revised: Attention ( QW Qi , KW Ki , V W Vi )

Original: Where W Qi , W Ki , and W Vi are learnable weight matrices to embed the queries, keys, and values.
Revised: Where W Qi , W Ki , and W Vi are learnable weight matrices to embed the queries, keys, and values.

Original: B D OMAIN DETAILS B. C LASSIC POMDP S
Revised: B Domain details B.1 Classic POMDPs

Original: We use Hallway ((Littman et al., 1995)) and HeavenHell ((Bonet, 1998)) from the classic POMDP literature.
Revised: We use Hallway ((Littman et al., 1995)) and HeavenHell ((Bonet, 1998)) from the classic POMDP literature.

Original: Hallway is a hallway gridworld with four rooms to the hallway’s south.
Revised: Hallway is a hallway gridworld with four rooms to the hallway’s south.

Original: The agent’s goal is to navigate to the fourth southern room despite very stochastic transition and observation dynamics.
Revised: The agent’s goal is to navigate to the fourth southern room despite very stochastic transition and observation dynamics.

Original: Due to the stochasticity, an agent must take several localizing actions before it can be certain of its surroundings.
Revised: Due to the stochasticity, an agent must take several localizing actions before it can be certain of its surroundings.

Original: Hallway’s observation space is an integer representing which walls the agent can see in its current cell (e.g. if the agent is at the end of the hallway, it would see the wall to its left, its right, and in front of it), and there are 5 actions: no-op, move forward, turn right, turn left, and turn around.
Revised: Hallway’s observation space is an integer representing which walls the agent can see in its current cell (e.g. if the agent is at the end of the hallway, it would see the wall to its left, its right, and in front of it), and there are 5 actions: no-op, move forward, turn right, turn left, and turn around.

Original: HeavenHell consists of a T-shaped grid with an oracle priest at the southern end, heaven at one end of the northern fork, and hell at the other end.
Revised: HeavenHell consists of a T-shaped grid with an oracle priest at the southern end, heaven at one end of the northern fork, and hell at the other end.

Original: On each episode reset, the location of heaven and hell may be swapped, and the agent can only learn heaven’s location by visiting the priest.
Revised: On each episode reset, the location of heaven and hell may be swapped, and the agent can only learn heaven’s location by visiting the priest.

Original: The optimal agent first navigates south, away from its goal, to the priest, then uses the priest’s observation to go to heaven.
Revised: The optimal agent first navigates south, away from its goal, to the priest, then uses the priest’s observation to go to heaven.

Original: HeavenHell’s observation space is an integer, representing the agent’s position in the world unless the agent is visiting the priest, in which case it recieves an indicator of whether heaven is on the left or the right side of the fork.
Revised: HeavenHell’s observation space is an integer, representing the agent’s position in the world unless the agent is visiting the priest, in which case it recieves an indicator of whether heaven is on the left or the right side of the fork.

Original: A HeavenHell agent can take one of 4 actions: move north, south, east, and west.
Revised: A HeavenHell agent can take one of 4 actions: move north, south, east, and west.

Original: In Hallway, the agent receives a living reward of 0, and a reward of 1 only after it reaches the goal state.
Revised: In Hallway, the agent receives a living reward of 0, and a reward of 1 only after it reaches the goal state.

Original: In HeavenHell, the agent receives a reward of 1 for reaching heaven, and -1 for reaching hell.
Revised: In HeavenHell, the agent receives a reward of 1 for reaching heaven, and -1 for reaching hell.

Original: B. G YM - GRIDVERSE
Revised: B.2 Gym-gridverse

Original: Gym-Gridverse ((Baisero & Katt, 2021)) contains a set of gridworlds featuring challenging and partially observable tasks.
Revised: Gym- Gridverse ((Baisero & Katt, 2021)) contains a set of gridworlds featuring challenging and partially observable tasks.

Original: We evaluate DTQN and our baselines on Gridverse’s memory and memoryfour-rooms environments.
Revised: We evaluate DTQN and our baselines on Gridverse’s memory and memoryfour-rooms environments.

Original: The agent’s observations consist of a limited forward view based on the agent’s position and orientation.
Revised: The agent’s observations consist of a limited forward view based on the agent’s position and orientation.

Original: In our experiments, the observation space is a vector of lengthintegers (representing the 2 × 3 grid as seen in Figure 3), and there are 6 actions: move forward, move backward, move left, move right, turn left, and turn right.
Revised: In our experiments, the observation space is a vector of lengthintegers (representing the 2 × 3 grid as seen in Figure 3), and there are 6 actions: move forward, move backward, move left, move right, turn left, and turn right.

Original: Memory consists of a central hallway with the southern end containing identical colored beacon tiles while the northern end holds two differently colored flags – one representing heaven and one representing hell.
Revised: Memory consists of a central hallway with the southern end containing identical colored beacon tiles while the northern end holds two differently colored flags – one representing heaven and one representing hell.

Original: Memory four rooms is a procedurally generated gridworld with a similar goal to Memory in that the agent must first find the colored beacon to know which colored flag to reach.
Revised: Memory four rooms is a procedurally generated gridworld with a similar goal to Memory in that the agent must first find the colored beacon to know which colored flag to reach.

Original: However, memory four rooms is much harder than memory because the layout of the grid as well as the beacon and flag locations generated randomly on each episode reset.
Revised: However, memory four rooms is much harder than memory because the layout of the grid as well as the beacon and flag locations generated randomly on each episode reset.

Original: The agent must gather information and explore its surroundings until it knows where the color of the beacon and the location of the flags.
Revised: The agent must gather information and explore its surroundings until it knows where the color of the beacon and the location of the flags.

Original: B. M EMORY CARDS
Revised: B.3 Memory cards

Original: Memory cards is a new domain based on the children’s card game Memory.
Revised: Memory cards is a new domain based on the children’s card game Memory.

Original: 5 pairs of cards are randomly shuffled and their locations are hidden.
Revised: 5 pairs of cards are randomly shuffled and their locations are hidden.

Original: At each time step, the agent sees the position of one random card (i.e. it is flipped face-up) and must guess the position of that card’s pair.
Revised: At each time step, the agent sees the position of one random card (i.e. it is flipped face-up) and must guess the position of that card’s pair.

Original: If guessed correctly, the pair of cards are removed from play; otherwise, the card is turned back face-down, and a new random card is revealed.
Revised: If guessed correctly, the pair of cards are removed from play; otherwise, the card is turned back face-down, and a new random card is revealed.

Original: This process continues until all N pairs have been removed from play.
Revised: This process continues until all N pairs have been removed from play.

Original: Each correct guess rewards the agent with a reward of 0, each incorrect guess a reward of -1.
Revised: Each correct guess rewards the agent with a reward of 0, each incorrect guess a reward of -1.

Original: The observation space is a vector of 10 integers, where each integer denotes the card at that position is hidden, face-up (in which case it shows that card’s value), or removed from play.
Revised: The observation space is a vector of 10 integers, where each integer denotes the card at that position is hidden, face-up (in which case it shows that card’s value), or removed from play.

Original: The agent’s action is the index for which card it thinks matches the currently revealed card.
Revised: The agent’s action is the index for which card it thinks matches the currently revealed card.

Original: Given enough time, a random policy can solve this task.
Revised: Given enough time, a random policy can solve this task.

Original: However, policies that can effectively remember their history and reason about unknown positions will perform best.
Revised: However, policies that can effectively remember their history and reason about unknown positions will perform best.

Original: B. C AR FLAG
Revised: B.4 Car flag

Original: Car flag tasks a car with driving across on a 1D line to the correct flag.
Revised: Car flag tasks a car with driving across on a 1D line to the correct flag.

Original: The car must first drive to the oracle flag and then to the correct endpoint.
Revised: The car must first drive to the oracle flag and then to the correct endpoint.

Original: The agent observation is a vector of 3 floats, including its position on the line, its velocity at each timestep, and, when it is at the oracle flag, it is also informed of the goal flag’s location.
Revised: The agent observation is a vector of 3 floats, including its position on the line, its velocity at each timestep, and, when it is at the oracle flag, it is also informed of the goal flag’s location.

Original: The agent’s action alters its velocity; it may accelerate left, perform a no-op (i.e. maintain current velocity), or accelerate right.
Revised: The agent’s action alters its velocity; it may accelerate left, perform a no-op (i.e. maintain current velocity), or accelerate right.

Original: The agent receives a reward offor reaching the goal flag, a reward of -1 for reaching the incorrect falg, and 0 otherwise.
Revised: The agent receives a reward offor reaching the goal flag, a reward of -1 for reaching the incorrect falg, and 0 otherwise.

Original: C A DDITIONAL LEARNING CURVES C. A DDITIONAL BASELINES
Revised: C Additional learning curves C.1 Additional baselines

Original: In addition to DRQN and DQN, we also compare our method to Action-based DRQN (ADRQN) (Zhu et al., 2017) and Deep Attention Recurrent Q-Network (DARQN)
Revised: In addition to DRQN and DQN, we also compare our method to Action-based DRQN (ADRQN) (Zhu et al., 2017) and Deep Attention Recurrent Q-Network (DARQN)

Original: (Sorokin et al., 2015).
Revised: (Sorokin et al., 2015).

Original: ADRQN uses the same structure as DRQN, but conditions the Q-function on previous actions as well as observations.
Revised: ADRQN uses the same structure as DRQN, but conditions the Q-function on previous actions as well as observations.

Original: DARQN similarly uses the structure of DRQN, but applies attention at each step between the LSTM’s last hidden state and the current observation.
Revised: DARQN similarly uses the structure of DRQN, but applies attention at each step between the LSTM’s last hidden state and the current observation.

Original: The results for these comparisons are shown in Figure 5.
Revised: The results for these comparisons are shown in Figure 5.

Original: In all cases, DTQN outperforms the baselines.
Revised: In all cases, DTQN outperforms the baselines.

Original: ADRQN performed better as a baseline than DRQN in gridverse memory 7x7 and hallway, but still performed worse than DTQN in terms of final success rate.
Revised: ADRQN performed better as a baseline than DRQN in gridverse memory 7x7 and hallway, but still performed worse than DTQN in terms of final success rate.

Original: In the Hallway domain, ADRQN and DARQN achieve high performance early, but DTQN achieves the best final success rate.
Revised: In the Hallway domain, ADRQN and DARQN achieve high performance early, but DTQN achieves the best final success rate.

Original: Unfortunately, we were unable to run these baselines on all our domains due to compute limitations, as DARQN took 12 hours per one million timesteps (compared to DTQN’s four hours per one million timesteps).
Revised: Unfortunately, we were unable to run these baselines on all our domains due to compute limitations, as DARQN took 12 hours per one million timesteps (compared to DTQN’s four hours per one million timesteps).

Original: [Figure]
Revised: [Figure]

Original: C. A BLATIONS
Revised: C.2 Ablations

Original: In this section, we plot the learning curves for the ablations in Table 1.
Revised: In this section, we plot the learning curves for the ablations in Table 1.

Original: Figure 6 refers to the “Transformer structure” section of Table 1.
Revised: Figure 6 refers to the “Transformer structure” section of Table 1.

Original: DTQN outperforms or performs competitively with all ablated versions.
Revised: DTQN outperforms or performs competitively with all ablated versions.

Original: The learning curves show both the ablated version of DTQN with GRU-like combine step and identity map reordering as well as the version containing only identity map reordering perform much worse in the Hallway domain compared to the original DTQN and the version of DTQN with only GRU-like gating.
Revised: The learning curves show both the ablated version of DTQN with GRU-like combine step and identity map reordering as well as the version containing only identity map reordering perform much worse in the Hallway domain compared to the original DTQN and the version of DTQN with only GRU-like gating.

Original: In gridverse memory 7x7, the identity map reordering version of DTQN learns slower than the other versions.
Revised: In gridverse memory 7x7, the identity map reordering version of DTQN learns slower than the other versions.

Original: Figure 7 refers to the “Positional encodings” section of Table 1.
Revised: Figure 7 refers to the “Positional encodings” section of Table 1.

Original: Again, DTQN performs better than or is similar to the other forms of positional encodings.
Revised: Again, DTQN performs better than or is similar to the other forms of positional encodings.

Original: Particularly, the memory cards domain shows the benefit of learned positional encodings, as our version clearly performs the best compared to the sinusoidal encodings and the variant of DTQN with no positional encodings.
Revised: Particularly, the memory cards domain shows the benefit of learned positional encodings, as our version clearly performs the best compared to the sinusoidal encodings and the variant of DTQN with no positional encodings.

Original: For more discussion of positional encodings, refer to Appendix E.
Revised: For more discussion of positional encodings, refer to Appendix E.

Original: Figure 8 refers to the “Intermediate Q-value prediction” section of Table 1.
Revised: Figure 8 refers to the “Intermediate Q-value prediction” section of Table 1.

Original: In all cases, DTQN outperforms or is competitive with our ablated versions, showing the important of the intermediate
Revised: In all cases, DTQN outperforms or is competitive with our ablated versions, showing the important of the intermediate

Original: Q-value prediction training regime.
Revised: Q-value prediction training regime.

Original: For more discussion of these results, refer to Section 5.5.
Revised: For more discussion of these results, refer to Section 5.5.

Original: [Figure]
Revised: [Figure]

Original: C. W ALL -C LOCK T IME C OMPARISON
Revised: C.3 Wall-Clock Time Comparison

Original: We also compare DTQN, DRQN, and DQN in terms of model size and speed in the Gridverse
Revised: We also compare DTQN, DRQN, and DQN in terms of model size and speed in the Gridverse

Original: Memory 5x5 domain.
Revised: Memory 5x5 domain.

Original: Table 3 shows the model size in terms of number of parameters as well as the amount of time required to reach 1M timesteps in the environment.
Revised: Table 3 shows the model size in terms of number of parameters as well as the amount of time required to reach 1M timesteps in the environment.

Original: DTQN contains significantly more parameters than DRQN and yet reaches 1M timesteps about 5 minutes slower than DRQN.
Revised: DTQN contains significantly more parameters than DRQN and yet reaches 1M timesteps about 5 minutes slower than DRQN.

Original: We believe this is because DTQN and the transformer are able to utilize the GPU hardware more effectively than the LSTM in DRQN.
Revised: We believe this is because DTQN and the transformer are able to utilize the GPU hardware more effectively than the LSTM in DRQN.

Original: Figure 9 shows the success rate of the models in the Gridverse
Revised: Figure 9 shows the success rate of the models in the Gridverse

Original: Memory 5x5 environment with respect to environment steps and wall-clock time.
Revised: Memory 5x5 environment with respect to environment steps and wall-clock time.

Original: DTQN is the most sample efficient in terms of both environment steps and wall-clock time.
Revised: DTQN is the most sample efficient in terms of both environment steps and wall-clock time.

Original: In Figure 9b, DTQN is the only agent capable of solving the task within the 50 minutes of wall-clock time.
Revised: In Figure 9b, DTQN is the only agent capable of solving the task within the 50 minutes of wall-clock time.

Original: To collect these results, we used a PC with Intel Core i7-8700K CPU @
Revised: To collect these results, we used a PC with Intel Core i7-8700K CPU @

Original: 3.70GHz x 12 and NVIDIA GeForce RTX 2070 Rev. A GPU.
Revised: 3.70GHz x 12 and NVIDIA GeForce RTX 2070 Rev. A GPU.

Original: Results may vary on other hardware.
Revised: Results may vary on other hardware.

Original: D A TTENTION VISUALIZATIONS
Revised: D Attention visualizations

Original: Figure 10 shows attention weights for a trajectory rolled out by DTQN in the gridverse memory 5x5 domain.
Revised: Figure 10 shows attention weights for a trajectory rolled out by DTQN in the gridverse memory 5x5 domain.

Original: The agent first looks for the colored beacons, then turns to navigate toward the flags.
Revised: The agent first looks for the colored beacons, then turns to navigate toward the flags.

Original: Crucially, in its second to last observation (magnified, left, in Figure 10), the agent can only see the green flag but needs to navigate to the blue flag.
Revised: Crucially, in its second to last observation (magnified, left, in Figure 10), the agent can only see the green flag but needs to navigate to the blue flag.

Original: It therefore attends to both the observation containing 
Revised: It therefore attends to both the observation containing the blue beacon as well as the initial observation containing the spatial relationship of the two flags

Original: [Figure]
Revised: [Figure]

Original: [Table]
Revised: [Table]

Original: the blue beacon as well as the initial observation containing the spatial relationship of the two flags (see magnified observations on right in Figure 10) and knows to move backwards towards the blue flag to achieve its goal.
Revised: (see magnified observations on right in Figure 10) and knows to move backwards towards the blue flag to achieve its goal.

Original: Similarly, Figure 11 displays attention weights for a trajectory rolled out by a trained DTQN agent in gridverse memory 7x7.
Revised: Similarly, Figure 11 displays attention weights for a trajectory rolled out by a trained DTQN agent in gridverse memory 7x7.

Original: The agent takes several actions at the beginning of the episode to localize itself.
Revised: The agent takes several actions at the beginning of the episode to localize itself.

Original: Notably, the agent encounters the yellow flag before seeing a beacon (magnified, top left in
Revised: Notably, the agent encounters the yellow flag before seeing a beacon (magnified, top left in

Original: Figure 11).
Revised: Figure 11).

Original: At this point, it does not know which flag is its goal, so it continues until it locates the red beacon (magnified, right).
Revised: At this point, it does not know which flag is its goal, so it continues until it locates the red beacon (magnified, right).

Original: This observation is attended to by all future observations, indicating the model understands this beacon dictates its goal for the episode.
Revised: This observation is attended to by all future observations, indicating the model understands this beacon dictates its goal for the episode.

Original: Finally, the agent finds the red flag (magnified, bottom left), and correctly navigates to it to successfully complete the episode.
Revised: Finally, the agent finds the red flag (magnified, bottom left), and correctly navigates to it to successfully complete the episode.

Original: [Figure] [Figure]
Revised: [Figure] [Figure]

Original: E P OSITIONAL ENCODINGS
Revised: E Positional encodings

Original: DTQN learns positional encodings, which gives the network information regarding the temporal position of each observation in its history.
Revised: DTQN learns positional encodings, which gives the network information regarding the temporal position of each observation in its history.

Original: In Figure 12, we examine the positional encodings
Revised: In Figure 12, we examine the positional encodings

Original: DTQN learns in various domains and compare them to the sinusoidal positional encodings from the original transformer (Vaswani et al., 2017).
Revised: DTQN learns in various domains and compare them to the sinusoidal positional encodings from the original transformer (Vaswani et al., 2017).

Original: DTQN learns unique positional encoding structures to match its domain.
Revised: DTQN learns unique positional encoding structures to match its domain.

Original: The positional encodings for the gridverse memory domain show high similarity scores, especially for positions 20 through 50, indicating the encodings may not be very valuable for this domain.
Revised: The positional encodings for the gridverse memory domain show high similarity scores, especially for positions 20 through 50, indicating the encodings may not be very valuable for this domain.

Original: Indeed, this notion is supported in Table 1 and Figure 7, where we see very similar performance between the DTQN agent with learned positional encodings and the variant without positional encodings.
Revised: Indeed, this notion is supported in Table 1 and Figure 7, where we see very similar performance between the DTQN agent with learned positional encodings and the variant without positional encodings.

Original: In contrast, DTQN’s learned encodings in the memory card domain are very distinct, which we can see in the low similarities everywhere except the symmetric diagonal in
Revised: In contrast, DTQN’s learned encodings in the memory card domain are very distinct, which we can see in the low similarities everywhere except the symmetric diagonal in

Original: Figure 12.
Revised: Figure 12.

Original: This indicates that these unique positional encodings are very valuable in this domain, which is supported by the poor performance of DTQN without any positional encodings on this domain (results shown in Table 1 and Figure 7).
Revised: This indicates that these unique positional encodings are very valuable in this domain, which is supported by the poor performance of DTQN without any positional encodings on this domain (results shown in Table 1 and Figure 7).

Original: We value the flexibility of learning unique positional encodings for each domain as it lets our model adapt to each environment.
Revised: We value the flexibility of learning unique positional encodings for each domain as it lets our model adapt to each environment.