SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Note that  $f_{\theta}^{1}$  and  $f_{\theta}^{2}$  form the current error at time  $t$  (i.e. the P term), and  $f_{\theta}^{3}$  forms the change in error (i.e. the D term). For accumulation strategies, using exponential smoothing with  $\alpha = 1$  for the first and third heads and a summation head for the second head will recover the P, I, and D terms for heads 1, 2, and 3, respectively. Note that the above assumes  $dt = 1$ , but the linear projections can be adjusted to take different  $dt$  values into account.
Code Release All code for implementations are provided in the supplemental material along with instructions for how to run experiments. The code can also be found at https://github.com/IanChar/GPIDE. The only experiment that cannot be run are the "real" cases for tokamak control.
Architecture We use the same general architecture for each of the RL methods in this paper (see Figure 5). Each input to the history encoders, policy functions, and  $Q$ -value functions have corresponding encoders. This setup closely follows what was done in Ni et al. [50]. The encoders are simply linear projections; however, in the case of our GRU history encoder we do linear projections followed by a ReLU activation (as done in Ni et al. [50]). Although hypothetically the policy only needs to take in history encoding,  $z_{t}$ , since int includes the current observation, we found it essential for the current observation to be passed in independently and have its own encoder.
Figure 5: General Policy and Q Function Architectures. This architecture is heavily inspired by Ni et al. [50]. The gray box shows the history encoder modules, and this is the only thing that changes between baseline methods in the tracking problems. Note that there are two encoders: one for the policy function and one for the  $Q$  value function. The purple boxes show the input encoders, and hyperparameters for these can be found in Table 6. We found the shortcut encoders to be essential to good performance. The architecture when using GRU is nearly identical; however, there is no "Transition Encoder" since Ni et al. [50] encodes  $(o_i, a_{i-1}, r_{i-1})$  for each time step instead.
B.1 GPIDE Implementation Details
In addition to what is mentioned in Section 3, we found that there were several choices that helped with training. First, there may be some scaling issues because  $o_{t} - o_{t-1}$  may be small or the result of summation type heads may result in large encodings. To account for this, we use batch normalization layers [34] before each input encoding and after each  $\ell^{h}$ .
There are very few nonlinear components of GPIDE. The only one that remains constant across all experiments is that a tanh activation is used for the final output of the encoder. For tracking tasks, the decoder  $g_{\theta}$  has 1 hidden layer with 64 units and uses a ReLU activation function. For PyBullet tasks,  $g_{\theta}$  is a linear function.
B.2 Recurrent and Transformer Baseline Details
Recurrent Encoder. For the recurrent encoder, we tried to match as many details as Ni et al. [50] as possible. We double checked our implementation against theirs and confirmed that it achieves similar performance.
Transformer Encoder. We follow the GPT2 architecture [58] for inspiration, and particularly the code provided in Karpathy [38]. In particular, we use a number of multi-headed self-attention blocks in sequence with residual connections. We use layer normalization [6] before multi-headed attention and out projections; however, we do not use dropout. The out projection for each multi-headed self-attention block has one hidden layer with four times the number of units as the embedding dimension. Although Melo [46] suggests using T-Fixup weight initialization, we found that more reliably high performance was achieved with the weight initialization of Radford et al. [58]. Lastly, we used the same representation for the history as GPIDE, i.e.  $(o_{t - 1},a_{t - 1},r_{t - 1},o_t - o_{t - 1})$ , since it results in better performance.
B.3 PID Baseline
To tune our PID baseline, we used Bayesian Optimization over the three (for SISO) or six (for MIMO) dimensional space. Specifically we use the library provided by Nogueira [51]. The output of the blackbox optimization is the average over 100 different settings (independent from the 100 settings used for testing). We allow the optimization procedure to collect as many samples as the RL methods. The final performance reported uses the PID controller with the best gains found during the
optimization procedure. The bounds for each of the tracking tasks were eyeballed to be appropriate, which potentially preferably skews performance.
C Hyperparameters
Because of resource restrictions, we were unable to do full hyperparameter tuning for each benchmark presented in this paper. Instead, we focused on ensuring that all history encoding methods were roughly comparable, e.g. dimension of encoding, number of parameters, etc. Tables 5 show selected hyperparameters, and the following subsections describe how an important subset of these hyperparameters were picked. Any tuning that was done was over three seeds using 100 fixed settings (different from the 100 settings used for testing).
Task Type	Learning Rate	Batch Size	Discount Factor	Policy Network	Q Network	Path Length Encoding
Tracking	3e-4	32 (256 for PIDE)	0.95	[24]	[256, 256]	100
PyBullet	3e-4	32 (256 for PIDE)	0.99	[256, 256]	[256, 256]	64
Table 5: SAC Hyperparameters. The "Path Length Encoding" is the amount of history each encoder gets to observe besides PIDE which, because of the nature of it, uses the entire episode.
Observation	Action	Reward	Transition	Policy Shortcut	Q Shortcut	History Encoding
GPIDE (Tracking)	8	N/A	N/A	8	8	64	64
GRU (Tracking)	8	N/A	N/A	N/A	8	64	64
Transformer (Tracking)	16	N/A	N/A	16	8	64	64
GPIDE (PyBullet)	32	16	16	64	8	64	128
Transformer (PyBullet)	48	16	16	48	8	64	128
Table 6: Dimension for the Input Encoders and Final History Encoding. The input encoders correspond to the output dimensions of the purple boxes in Figure 5. By "History Encoding" size we mean the dimension of  $z_{t}$ .
Task Type	D	gθHidden Size
Tracking	16	[64]
PyBullet	32	[]
Table 7: GPIDE Specific Hyperparameters. Recall that  $\overline{D}$  corresponds to the output dimension of  $f_{\theta}$ . Empty brackets for the hidden size means that  $g_{\theta}$  is a linear function.
C.1 Hyperparameters for Tracking Tasks
For tracking tasks, we tried using a history encoding size of 32 and 64 for GRU, and we found that performance was better with 64. This is surprising since PIDE can perform well in these environments even though its history encoding is much smaller (3 or 6 dimensional). To make it a fair comparison, we set the history encoding dimension for GPIDE and transformer to be 64 as well. We use one layer for GRU. For the transformer-specific hyperparameters we choose half of what appears in the PyBullet tasks.
C.2 Hyperparameters for PyBullet Task
For the PyBullet tasks, we simply tried to emulate most of the hyperparameters found in Ni et al. [50]. For the transformer, we choose to use similar hyperparameters to those found in Melo [46]. However, we found that, unlike the tracking tasks, positional encoding hurts performance. As such, we do not include it for PyBullet experiments.
C.3 Hyperparameters for Ablations
For the ablations of GPIDE, we use  $\alpha = 0.01, 0.1, 0.25, 0.5, 0.9, 1.0$  for the smoothing parameters when only exponential smoothing is used. When using exponential smoothing and summation, the  $\alpha = 0.01$  head is replaced with a summation head. The attention version of GPIDE replaces all six of these heads with attention.
Task Type	Number of Layers	Number of Heads	Embedding Size per Head
Tracking	2	4	8
PyBullet	4	8	16
Table 8: Transformer Specific Hyperparameters
Encoder	SISO Tracking	MIMO Tracking (2D)	PyBullet
Transformer	25,542	25,644	793,868-795,026
GRU	14,240	14,264	74,816-75,248
GPIE	13,228	13,288	75,296-76,486
GPIE-ES	12,204	12,264	50,720-51,910
GPIE-ESS	12,204	12,264	50,720-51,910
GPIE-Attention	15,276	15,336	99,872-101,062
Table 9: Number of Parameters in History Encoder Modules. The number of parameters corresponds to the gray boxes in Figure 5. The difference in SISO vs MIMO and the PyBullet tasks is due to the different observation and action space dimensionalities.
D Computation Details
We used an internal cluster of machines to run these experiments. We mostly leveraged Nvidia Titan X GPUs for this, but also used a few Nvidia GTX 1080s. It is difficult to get an accurate estimate of run time since job loads vary drastically on our cluster from other users. However, to train a single policy on DMSD to completion (1 million transitions collected, or 1,000 epochs) using PIDE takes roughly 4.5 hours, using GPIDE takes roughly 17.25 hours, using a GRU takes roughly 14.5 hours, and using a transformer takes roughly 21 hours. This is similar for other tracking tasks. For PyBullet tasks, using GPIDE took roughly 43.2 hours and using a transformer took roughly 64.2 hours. We note that our implementation of GPIDE is somewhat naive and could be vastly improved. In particular, for exponential smoothing and summation heads,  $w_{t}$  can be cached to save on compute, which is not being done currently. This is a big advantage in efficiency that GPIDE (especially one without attention heads) has over transformers.
E Environment Descriptions
E.1 Mass Spring Damper
DMSD
Figure 6: Diagram of the Mass Spring Damper Environments. The diagram on the left the Mass Spring Damper (MSD) environment, and the diagram on the right shows the Double Mass Spring Damper (DMSD) environment. In the diagram, we have labelled the system parameters and the parts of the observation. The dotted line shows where the center of the mass is located with no force applied, and the current position of the mass is measured with respect to this point.
For both MSD and DMSD, the observations include the current mass position(s), the target reference position(s), and the last action played. Each episode lasts for 100 time steps. For all RL methods, the action is a difference in force applied to the mass, but for the PID the action is simply the force
System Parameter	Fixed	Small	Large
Damping Constant	U(4.0, 4.0)	U(3.5, 5.5)	U(2.0, 10.0)
Spring Constant	U(2.0, 2.0)	U(1.75, 3.0)	U(0.5, 6.0)
Mass	U(20.0, 20.0)	U(17.5, 40.0)	U(10.0, 100.0)
Table 10: MSD and DMSD System Parameter Distributions. Each episode system parameters are uniformly at random drawn from these bounds.
System Parameter	No Friction	Friction
Total Friction	U(0.0, 0.0)	U(0.05, 0.25)
Static Friction (Proportion)	U(0.0, 0.0)	U(0.25, 0.75)
Mass	U(15.0, 25.0)	U(5.0, 35.0)
Table 11: Navigation System Parameter Distributions. Each episode system parameters are uniformly at random drawn from these bounds. The static friction parameter drawn is the proportion of the total friction that is static friction.
to be applied to the mass at that time. The force is bounded between -10 and  $10N$  for MSD and -30 and  $30N$  for DMSD. Each episode, system parameters are drawn from a uniform distribution with bounds shown in Table 10 (they are the same for both MSD and DMSD). Targets are drawn to uniformly at random to be -1.5 to  $1.5m$  offset from the masses' resting positions.
E.2 Navigation Environment
Like the MSD and DMSD environments, the navigation experiment lasts 100 time steps each episode. Additionally, the observation includes position signal, target locations, and the last action. For all methods we set the action to be the change in force, and the total amount of force is bounded between -10 and  $10N$ . The penalty on the reward is equal to 0.01 times the magnitude of the change in force. In addition, the maximum magnitude of the velocity for the agent is bounded by  $1.0m / s$ . The agent always starts at the location  $(0,0)$ , and the target is picked uniformly at random to be within a box of length 10 centered around the origin.
Every episode, the mass, kinetic friction coefficient, and static friction coefficient is sampled. The friction is sampled by first sampling the total amount of friction in the system, and then sampling what proportion of the total friction is static friction. All distributions for the system parameters are uniform, and we show the bounds in Table 11.
E.3 Tokamak Control Environment
Simulator Our simulator version of the tokamak control is inspired by equations used by Boyer et al. [9], Scoville et al. [63]. In particular, we use the following relations for stored energy,  $E$ , and rotation,  $v_{\mathrm{rot}}$ :
$$
\dot {E} = P - \frac {E}{\tau_ {E}}
$$
$$
\dot {v} _ {\mathrm {r o t}} = C _ {\mathrm {r o t}} T - \frac {v _ {\mathrm {r o t}}}{\tau_ {m}}
$$
where  $P$  is the total power,  $T$  is the total torque,  $\tau_{E}$  is the energy confinement time,  $\tau_{m}$  is the momentum confinement time, and  $C_{\mathrm{rot}}$  is a quantity relying on the ion density and major radius of the plasma. We treat  $\tau_{m}$  and  $C_{\mathrm{rot}}$  as constants with values of 0.1 and 80.0 respectively.
We base the energy confinement time off of the ITERH-98 scaling [68]. This uses many measurements of the plasma, but we focus on a subset of these and treat the rest as constants. In particular,
$$
\tau_ {E} = C _ {E} I ^ {0. 9 5} B ^ {0. 1 5} P ^ {- 0. 6 9}
$$
where  $C_E$  is a constant value we set to be 200,  $I$  is the plasma current, and  $B$  is the toroidal magnetic field. To relate the stored energy to  $\beta_N$  we use the rough approximation
$$
\beta_ {N} = C _ {\beta} \left(\frac {a B}{I}\right) E
$$
Minor Radius (m)	Plasma Current (MA)	Toroidal Magnetic Field (T)
N(0.589, 0.02)	N(1e6, 1e5)	N(2.75, 0.1)
Table 12: Tokamak Control Simulator Distributions.
Figure 7: Power and Torque Bounds. The region outlined in blue shows the possible power-torque configurations. The dots show possible requests, and the corresponding red  $X$  marks show the actual achieved power-torque setting.
where  $C_{\beta}$  is a constant we set to be 5, and  $a$  is the minor radius of the plasma. For  $a, I$ , and  $B$ , we sample these from the distribution described in Table 12 for each episode. Lastly, we add momentum to the stored energy. That is, the stored energy derivative at time  $t$ ,  $\dot{E}_t$ , is
$$
\dot {E} _ {t} = 0. 5 \left(P _ {t} - \frac {E _ {t}}{\tau_ {E}}\right) + 0. 5 \dot {E} _ {t - 1}
$$
The actions for all control methods is the amount of change for the power and torque. Because the total amount of power and torque injected rely on the beams, they are not totally disentangled. In Figure 7, we show the bounds for the action space. Furthermore, we bound the amount that power and torque can be changed by roughly  $40MW / s$  and  $35Nm / s$ , respectively. Each step is 0.025 seconds.
Each episode lasts for 100 increments of 0.025 seconds. The observations are the current  $\beta_{N}$  and rotation values, their reference values, and the current power and torque settings. We make the initial  $\beta_{N}$  and rotation relatively small in order to simulate the plasma ramping up. We let the  $\beta_{N}$  and rotation targets be distributed as  $\mathcal{U}(1.75,2.75)$  and  $\mathcal{U}(25.0,50.0)$  rad/s, respectively.
"Real" For the real versions of the tokamak control experiments, most of the previous (such as action bounds and target distributions) stays the same. This data-driven simulator is based on the one from Char et al. [14], and we refer the reader there for more details. That being said, there are some differences between the architecture presented there and the one used in this work. Our network is a recurrent network that uses a GRU, has four hidden layers with 512 units each, and outputs the mean and log variance of a normal distribution describing how  $\beta_{N}$  and rotation will change. In addition to power and torque, it takes in measurements for the plasma current, the toroidal magnetic field, n1rms (a measurement related to the plasma' stability), and 13 other actuator requests for gas control and plasma shaping. In addition to sampling from the normal distribution outputted by the network, we train an ensemble of ten networks, and an ensemble member is selected every episode. We use five of
these models during training and the other five during testing. Along with an ensemble member being sampled each episode, we also sample a historical run, which determines the starting conditions of the plasma and how the other inputs to the neural network which are not modelled evolve over time. Recall that 100 fixed settings are used to evaluate the policy every epoch of training. In this case, a setting consists of targets, an ensemble member, and a historical run.
F.1 Experiments Using VIB + GRU
As shown in this work, using a GRU for a history encoder often results in a policy that is ill-equipped to handle changes in the dynamics not seen at train time. One may wonder whether using other robust RL techniques is able to mask this inadequacy of GRU. To test this, we look at adding Variational Information Bottlenecking (VIB) to our GRU baseline [4]. Previous works applying this concept to RL usually do not consider the same class of POMDPs as us [42, 33]; however, Eysenbach et al. [21] does have a baseline that uses VIB with a recurrent policy.
To use VIB with RL, we alter the policy network so that it encodes input to a latent random variable, and the decodes into an action. Following the notation of Lu et al. [42], let this latent random variable be  $Z$  and the random variable representing the input of the network be  $S$ . The goal is to learn a policy that maximizes  $J(\pi)$  subject to  $I(Z, S) \leq I_C$ , where  $I(Z, S)$  is the mutual information between  $Z$  and  $S$ , and  $I_C$  is some given threshold. In practice, we optimize the Lagrangian. Where  $\beta$  is a Lagrangian multiplier,  $p(Z|S)$  is the conditional density of  $Z$  outputted by the encoder, and  $q(Z)$  is the prior, the penalizer is  $-\beta \mathbb{E}_S[p(Z|S)||q(Z)]$ . Like other works, we assume that  $q(Z)$  is a standard multivariate normal.
We alter our GRU baseline for tracking tasks so that the policy uses VIB. This is not entirely straightforward since our policy network is already quite small. We choose to keep as close to original policy architecture as possible and set the dimension of the latent variable,  $Z$ , to be 24. Note that this change has no affect on the history encoder; this only affects the policy network. For our experiments, we set  $\beta = 0.1$ , but we note that we may be able to achieve better performance through more careful tuning or annealing of  $\beta$ .
In any case, we do see that VIB helps with robustness in many instances (see Table 13). However, the cases where there are improvements are instances where the GRU policy already did a good job at generalizing to the test environment. These are primarily the MSD and DMSD environments where the system parameters drawn during training time are simply a subset of those drawn during testing time (interestingly, this notion of dynamics generalization matches the set up of the experiments presented in Lu et al. [42]). Surprisingly, in the navigation and tokamak control experiments, where there are more complex differences between the train and test environments, VIB can sometimes hurt the final performance.
PID Controller	GRU	GRU+VIB	Transformer	PIDE	GPIDE
MSD Fixed / Fixed	-6.14 ± 0.02	-5.76 ± 0.02	-5.73 ± 0.01	-5.75 ± 0.01	-5.69 ± 0.00	-5.76 ± 0.01
MSD Fixed / Large	-11.39 ± 0.09	-12.52 ± 0.11	-12.50 ± 0.14	-10.87 ± 0.05	-11.44 ± 0.03	-11.61 ± 0.07
MSD Small / Small	-7.49 ± 0.03	-7.02 ± 0.01	-7.01 ± 0.01	-7.15 ± 0.02	-7.14 ± 0.01	-7.12 ± 0.04
MSD Small / Large	-11.18 ± 0.09	-9.82 ± 0.07	-9.57 ± 0.03	-10.01 ± 0.03	-10.88 ± 0.04	-10.43 ± 0.14
DMSD Fixed / Fixed	-15.33 ± 0.14	-16.20 ± 0.31	-15.83 ± 0.28	-15.41 ± 0.13	-12.64 ± 0.04	-13.49 ± 0.22
DMSD Fixed / Large	-27.59 ± 0.44	-37.21 ± 0.35	-35.34 ± 0.28	-28.16 ± 0.17	-25.29 ± 0.18	-27.54 ± 0.33
DMSD Small / Small	-21.78 ± 0.14	-22.49 ± 0.34	-22.51 ± 0.24	-20.56 ± 0.16	-18.09 ± 0.04	-18.67 ± 0.17
DMSD Small / Large	-26.57 ± 0.22	-31.27 ± 0.36	-30.93 ± 0.34	-26.04 ± 0.24	-23.82 ± 0.13	-23.65 ± 0.20
Nav Sim / Sim	-17.23 ± 0.18	-13.82 ± 0.01	-14.69 ± 0.02	-13.68 ± 0.01	-13.74 ± 0.00	-13.65 ± 0.00
Nav Sim / Real	-23.87 ± 0.29	-29.85 ± 0.55	-39.57 ± 0.24	-22.84 ± 0.11	-20.37 ± 0.08	-21.23 ± 0.12
βN Sim / Sim	-8.09 ± 0.00	-7.19 ± 0.00	-7.24 ± 0.01	-7.22 ± 0.00	-8.71 ± 0.02	-7.21 ± 0.01
βN Sim / Real	-16.41 ± 0.30	-31.21 ± 1.67	-32.19 ± 1.19	-31.49 ± 3.66	-43.78 ± 6.46	-26.83 ± 1.36
βN-Rotation Sim / Sim	-27.56 ± 0.08	-18.53 ± 0.02	-18.61 ± 0.12	-18.79 ± 0.09	-21.36 ± 0.05	-18.45 ± 0.03
βN-Rotation Sim / Real	-30.08 ± 0.95	-45.91 ± 2.10	-44.24 ± 1.33	-48.23 ± 0.25	-60.23 ± 3.20	-41.86 ± 0.69
Average	-18.33	-20.12	-21.14	-18.71	-19.58	-17.51
Table 13: Tracking Experiments with GRU+VIB. We use green and red text to highlight significant improvements and deteriorations in performance over vanilla GRU. We only highlight a subset of configurations since we are focused on the robustness properties. This table shows average (unnormized) returns.
F.2 Lookback Size Ablations
To better understand the role of the maximum lookback size (i.e. the amount of history used to form the encoding) of GPIDE, we repeat the PyBullet experiments using a lookback size of 4, 16, 64, and 128 with and without attention (labelled GPIDE and GPIDE-ESS, respectively). Figures 8 and 9 show performance curves for GPIDE and GPIDE-ESS respectively. It is clear that there is a massive increase in improvement when expanding the maximum size of the lookback from 4 to 16. For the most part, this trend continues expanding the lookback size from 16 to 64; however, it seems pushing from 64 to 128 yields mixed results. For some tasks, such as HalfCheetah-P, expanding the lookback to 128 results in noticeable improvements both with and without attention. For other tasks, such as Hopper-V, this expansion yields slightly worse performance, possibly because of training stability issues.
One interesting observation is that, when attention is included, this decrease in performance from expanding lookback can occur when increasing from 16 to 64 (see HalfCheetah-V and Walker-V). At the same time, however, it appears the GPIDE can sometimes maintain good performance when expanding the lookback to 128 when GPIDE-ESS cannot (Walker-P).
Figure 8: GPIDE Lookback Ablation. Each curve shows the average over four seeds and the standard error of each.
Figure 9: GPIDE-ESS Lookback Ablation. Each curve shows the average over four seeds and the standard error of each.
G Further Results
In this Appendix, we give further evaluation of the evaluation procedure. In addition, we give full tables of results for normalized and unnormalized scores for all methods. We also show performance
traces. Note that the percentage changes in Table 4 do not necessarily reflect tables in this section since they report all combinations of environment variants.
G.1 Evaluation Procedure
As stated in the main paper, for tracking tasks, we fix 100 settings (each comprised of targets, start state, and system parameters) that are used to evaluate the policy for every epoch of training (i.e. for every epoch the evaluation returns is the average over all 100 settings returns). We use a separate 100 settings when tuning. For the final returns, we average over the last  $10\%$  of recorded evaluations.
For the PyBullet tasks, we use ten different rollouts for evaluation following Ni et al. [50]. We also average over the last  $20\%$  of recorded evaluations like they do.
Normalized Table Scores. We now give an in-depth explanation of how the scores in the table are computed. Let  $\pi_{(b,i)}$  be the policy trained with baseline method  $b$  (e.g. with GPIDE, transformer, or GRU encoder) on environment variant  $i$  (e.g. fixed, small, or large). Let  $J_{j}(\pi_{(b,i)})$  be the evaluation of policy  $\pi_{(b,i)}$  on environment variant  $j$ , i.e. the average returns over all seeds and episodes. The normalized score for policy  $\pi_{(b,i)}$  on variant  $j$  is then