SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
\nabla_ {\omega} L _ {\mathrm {s h a}} = \left(Q _ {\omega} \left(f _ {\phi} (h), a\right) - Q ^ {\operatorname {t a r}} \left(h, a, o ^ {\prime}, r\right)\right) \nabla_ {\omega} Q _ {\omega} \left(f _ {\phi} (h), a\right) - \underbrace {\nabla_ {\omega} Q _ {\omega} \left(f _ {\phi} (h) , \pi_ {\nu} \left(f _ {\phi} (h)\right)\right)} _ {\text {e x t r a g r a d i e n t t e r m}} \tag {37}
$$
Comparing the two gradient formulas across the two settings, we find the extra gradient terms using the same value of  $\phi$  for  $\phi_{\pi}$  and  $\phi_Q$  in  $L_{\mathrm{sep}}$ :
$$
\begin{array}{l} \nabla_ {\phi_ {\pi}} L _ {\operatorname {s e p}} (\phi , \phi , \omega , \nu) + \nabla_ {\phi_ {Q}} L _ {\operatorname {s e p}} (\phi , \phi , \omega , \nu) - \nabla_ {\phi} L _ {\operatorname {s h a}} (\phi , \omega , \nu) \\ = \nabla_ {z} Q _ {\omega} \left(f _ {\phi} (h), \pi_ {\nu} \left(f _ {\phi} (h)\right)\right) \nabla_ {\phi} f _ {\phi} (h) \tag {38} \\ \end{array}
$$
$$
\nabla_ {\omega} L _ {\operatorname {s e p}} (\phi , \phi , \omega , \nu) - \nabla_ {\omega} L _ {\operatorname {s h a}} (\phi , \omega , \nu) = \nabla_ {\omega} Q _ {\omega} \left(f _ {\phi} (h), \pi_ {\nu} \left(f _ {\phi} (h)\right)\right) \tag {39}
$$
To remove these extra gradient terms, we propose freezing the critic parameters  $\phi$  and  $\omega$  in the actor loss for the shared encoder setting, resulting in a modified loss function  $L_{\mathrm{sha-ours}}(\phi ,\omega ,\nu)$ :
$$
L _ {\text {s h a - o u r s}} (\phi , \omega , \nu) = L _ {Q} (\phi , \omega) - Q _ {\bar {\omega}} \left(f _ {\bar {\phi}} (h), \pi_ {\nu} \left(f _ {\phi} (h)\right)\right). \tag {40}
$$
This offers a theoretical explanation for the practice of detaching critic parameters in the actor loss, adopted in prior works (Yarats et al., 2021a,b; Yang and Nguyen, 2021). It is worth noting that these works take an additional step of freezing the encoder parameters present in actions, resulting in an expression like  $Q_{\overline{\omega}}(f_{\phi}(h), \pi_{\nu}(f_{\overline{\phi}}(h)))$ . In our preliminary experiments conducted on PyBullet tasks, we observed negligible differences in the empirical performance between their method and ours. Consequently, to maintain consistency with our theoretical insights, we have chosen to implement our approach  $L_{\mathrm{sha - ours}}$  in all of our experimental setups.
Implementation of sequence encoders in our experiments. We train both LSTM and Transformer sequence encoders with a full context length equal to the episode length on all tasks, except for PyBullet tasks where we follow POMDP-baselines to use a context length of 64. The training input of sequence encoders is based on a history of  $t$  observations and  $t$  actions (with the first action zero-padded). The history is embedded as follows according to the POMDP-baselines. For state-based tasks (T-Mazes and PyBullet), the observations and actions are individually embedded through a linear and ReLU layer. For pixel-based tasks (Passive Visual Match and Key to Door), the observations (images) are embedded through a small convolutional neural network. For all tasks, the sequence of embedded observations and actions are then concatenated along their final dimension to form the input sequence to the sequence model.
For LSTMs (Hochreiter and Schmidhuber, 1997), we train LSTMs with a hidden size of 128, varying the number of layers from 1, 2, 4. We find single-layer LSTM performs best in T-Mazes and Passive Visual Match, and two-layer best in Key-to-Door, which are reported in Fig. 3 and Fig. 6. See Table 3 for details.
For Transformers, we utilize the GPT-2 model (Radford et al., 2019) implemented by Hugging Face Transformer library (Wolf et al., 2019). Our Transformer is a stack of  $N$  layers with  $H$ -headed self-attention modules. It includes causal masks to condition only on past context. A sinusoidal positional encoding (Vaswani et al., 2017) is added to the embedded sequence. The same dropout rate is used for regularization on the embedding, residual, and self-attention layers. We tune  $(N,H)$  from  $(1,1),(2,2),(4,4)$ . We find  $(N,H) = (1,1)$  performs best in Passive T-Maze and Passive Visual Match, and  $(2,2)$  best in Active T-Maze and Key-to-Door, which are reported in Fig. 3 and Fig. 6. See Table 4 for details.
For PyBullet tasks, we follow POMDP-baselines to use single-layer LSTMs and Transformers.
Table 3: LSTM hyperparameters used for all experiments.
Hyperparameter	Value
State
embedding	Obs. embedding size	32
Act. embedding size	16
Pixel
embedding	No. channels	(8, 16)
Kernel size	2
Stride	1
Obs. embedding size	100
Act. embedding size	0
LSTM	Hidden size	128
No. layers	1, 2, 4
Table 4: Transformer hyperparameters used for all experiments.
Hyperparameter	Value
State
embedding	Obs. embedding size	64
Act. embedding size	64
Pixel
embedding	No. channels	(8, 16)
Kernel size	2
Stride	1
Obs. embedding size	100
Act. embedding size	0
GPT	Dropout	0.1
No. heads (H)	1, 2, 4
No. layers (N)	1, 2, 4
Implementation of RL algorithms in our experiments. We use SAC-Discrete (Christodoulou, 2019) for pixel-based tasks, and TD3 (Fujimoto et al., 2018) for PyBullet tasks, following POMDP-baselines. We use DDQN (van Hasselt et al., 2016) for T-Maze tasks, which outperforms SAC-Discrete in our preliminary experiments. We use epsilon-greedy exploration strategy in DDQN with a linear schedule, where the ending epsilon is  $\frac{1}{T}$  with  $T$  being the episode length. This ensures that the probability of always taking deterministic actions throughout an episode asymptotically approaches a constant:
$$
\lim  _ {T \rightarrow + \infty} (1 - \epsilon) ^ {T} = \lim  _ {T \rightarrow + \infty} \left(1 - \frac {1}{T}\right) ^ {T} = \frac {1}{e} \approx 0. 3 6 8 \tag {41}
$$
where  $e$  is the base of the natural logarithm. This approach is critical to solving T-Maze tasks which strictly require a deterministic policy. Table 5 summarizes the details of RL agents.
Table 5: RL agent hyperparameters used in all experiments.
Hyperparameter	Value
Network hidden size	(256, 256)
Discount factor (γ)	0.99
Target update rate	0.005
Replay buffer size	10^6
Learning rate	0.0003
Batch size	64
DDQN	epsilon greedy schedule	linear(1.0, 1/T, schedule_steps)
schedule_steps	0.1* num Episodes
SAC-Discrete	entropy temperature	0.1
Table 6: Training hyperparameters in our experiments.
Tasks	context_length	num Episodes
Passive T-Maze	1500	8000
Active T-Maze	500	4000
Passive Visual Match	1000	10000
Key-to-Door	500	4000
C.2 Additional Results
Fig. 11 shows the learning curves of training Transformers with varying numbers of layers and heads.
Similar to the scaling experiments on Transformers we demonstrated in Sec. 5.2, we conduct an ablation study on scaling the number of layers in LSTMs from 1, 2 to 4 in credit assignment tasks. Fig. 12 shows the results that multi-layer (stacked) LSTMs do not help performance, aligned with empirical findings in stacked LSTMs (Pascanu et al., 2014).
Figure 11: Learning curves of scaling the number of layers and heads in Transformers in Key-to-Door tasks, associated with Fig. 7.
Figure 12: Scaling the number of layers in LSTMs does not affect performance much.
C.3 Plotting Details
For all the learning curve plots in this paper, we use seaborn.lineplot (Waskom, 2021) to show the mean and its  $95\%$  confidence interval across 10 seeds. For the aggregation plots, we also use seaborn.lineplot to report the final performance, which is averaged over the evaluation results during the last  $5\%$  of interactions.
C.4 Training Compute
Here we report the training compute of the experiments in Sec. 5.1 and Sec. 5.2, because the PyBullet experiments (Sec. 5.3) are relatively cheap to compute. Each experiment run was carried out on a single A100 GPU and a single CPU core. For Transformer-based RL, the GPU memory usage is approximately proportional to the square of the context lengths, with a maximum usage of 4GB for Passive T-Maze with a context length of 1500. For LSTM-based RL, GPU memory usage is roughly linearly proportional.
In our tasks, the context length equals the episode length (and also memory length), thus the total training time is proportional to num Episodes * context_length**2 * update_frequency for both Transformers and LSTMs. The update Frequency, set as 0.25, is the ratio of parameter update w.r.t. environment step. Additionally, for multi-layer Transformers, the training time is roughly linear to the number of layers. Table 6 summarizes the training hyperparameters.
In Passive T-Maze with a memory length of 1500, it took around 6 and 4 days to train Transformer-based and LSTM-based RL, respectively. In Passive Visual Match with a memory length of 1000, it took around 5 days for both Transformer-based and LSTM-based RL.
Our work enhances understanding of long-term memory capability in Transformer-based RL. This improved memory capability, while offering advancements in RL, may pose privacy concerns if misused, potentially enabling systems to retain and misuse sensitive information. As this technology develops, strict data privacy measures are essential. However, negative impacts directly tied to our foundational research are speculative, as we propose no specific application of this technology.