SlowGuess/NeurIPS_2023_OCR huggingface dataset:
For the following tracking problems we use the Soft Actor Critic (SAC) [26] algorithm with each of the different methods for encoding observation history. Following Ni et al. [50], we make two separate instantiations of the encoders for the policy and value networks, respectively. Since the tracking problems are relatively simple, we use a small policy network consisting of 1 hidden layer with 24 units; however, we found that we still needed to use a relatively large Q network consisting of 2 hidden layers with 256 units each to solve the problems. All hyperparameters remain fixed across baselines and tracking tasks; only the history encoders change.
For the recurrent encoder, we use a GRU and follow the implementation of Ni et al. [50] closely. Our transformer encoder closely resembles the GPT2 architecture [58], and it also includes positional encodings for the observation history. For GPIDE, we use  $H = 6$  heads: one summation head, two attention heads, and three exponential smoothing heads (with  $\alpha = 0.25, 0.5, 1.0$ ). This choice was not optimized, but rather was picked so that all types of heads were included and so that GPIDE has roughly the same amount of parameters as our GRU baseline. As a reference point for these RL methods, we also evaluate the performance of a tuned PID controller. Not only do PID controllers have an incredibly small number of parameters compared to the other RL-based controllers, but the training procedure is also much more straightforward since it can be posed as a black-box optimization over the returns. While there exists many sophisticated extensions of the PID controller (especially in MIMO systems [8]), we only consider the vanilla PID controller since we believe it serves as a good reference point. All methods are built on top of the rlkit library [57]. More details about implementations, hyperparameters, and computation can be found in Appendices B, C, and D, respectively. We also include additional experiments regarding variational information bottlenecking (VIB) and lookback size ablations in Appendices F.1 and F.2. Implementations can be found at https://github.com/IanChar/GPIDE.
5.1 Tracking Problems
In this subsection we consider a number of tracking problems. For each environment, the observation consists of the current signals, the reference values, and additional information about the last action made. Unless stated otherwise, the reward is as described in Section 2. More information about environments can be found in Appendix E. To make a fair comparison against PID controls, we choose to only encode the history of observations. For evaluation, we use 100 fixed settings of the environment (each setting consists of targets and system parameters). To avoid overfitting to these 100 settings, we used a separate set of 100 settings and averaged over 3 seeds when developing our methods. We evaluate policies throughout training, but report the average over the last  $10\%$  of evaluations as the final returns. We allow each policy to collect one million environment transitions, and all scores are averaged over 5 seeds. Lastly, each table shows scores formed by scaling the returns by the best and worst average returns across all methods in a particular variant of the environment, where scores of 0 and 100 correspond to the worst and best returns respectively.
Environment (Train/Test)	PID Controller	GRU	Transformer	PIDE	GPIDE
MSD Fixed/Fixed	0.00 ± 3.96	83.73 ± 3.48	85.79 ± 1.98	100.00 ± 0.66	83.72 ± 2.86
MSD Small/Small	0.00 ± 5.58	100.00 ± 1.59	73.27 ± 4.98	75.51 ± 1.31	80.21 ± 8.59
MSD Fixed/Large	36.58 ± 2.86	0.00 ± 3.42	53.70 ± 1.71	34.92 ± 0.93	29.55 ± 2.32
MSD Small/Large	43.52 ± 2.82	87.63 ± 2.28	81.44 ± 0.82	53.21 ± 1.31	68.03 ± 4.43
MSD Large/Large	45.60 ± 1.71	100.00 ± 0.61	92.60 ± 1.49	69.88 ± 0.69	93.03 ± 1.27
Average	25.14	74.27	77.36	66.70	70.91
DMSD Fixed/Fixed	24.33 ± 3.97	0.00 ± 8.69	22.05 ± 3.58	100.00 ± 1.08	76.23 ± 6.26
DMSD Small/Small	16.17 ± 3.09	0.00 ± 7.79	43.74 ± 3.70	100.00 ± 0.94	86.74 ± 3.94
DMSD Fixed/Large	63.59 ± 2.91	0.00 ± 2.28	59.84 ± 1.13	78.77 ± 1.16	63.89 ± 2.16
DMSD Small/Large	70.35 ± 1.44	39.26 ± 2.37	73.81 ± 1.60	88.52 ± 0.83	89.66 ± 1.33
DMSD Large/Large	78.77 ± 1.97	52.01 ± 2.01	84.45 ± 1.41	86.90 ± 0.18	100.00 ± 0.91
Average	50.64	18.25	56.78	90.84	83.30
Total Average	37.89	46.26	67.07	78.77	77.11
Table 1: Mass Spring Damper Task Results. The scores presented are averaged over five seeds and we show the standard error for each score.
Mass Spring Damper Tracking The first tracking task is the control of a classic 1D toy physics system in which there is a mass attached to a wall by a spring and damper. The goal is then to apply a force to the mass in order to move it to a given reference location. There are three system parameters to consider here: the mass, spring constant, and damping factor. We also consider the substantially more difficult problem in which there are two masses sandwiched between two walls, and the masses are connected to the walls and each other by springs and dampers (see Appendix E.1 for a diagram of this). Overall there are eight system parameters (three spring constants, three damping factors, and two masses) and two actuators (a force applied to each mass). We refer to the first problem as Mass-Spring-Damper (MSD) and the second problem as Double-Mass-Spring-Damper (DMSD).
Additionally, we test how adaptive these policies are by changing system parameters in a MetaRL-type fashion (i.e. for each episode we randomly select system parameters and then fix them for the rest of the episode). Similar to Packer et al. [54], we train the policies on three versions of the environment: one with no variation in system parameters, one with a small amount of variation, and one with a large amount of variation. We evaluate all policies on the version of the environment with large system parameter variation to test generalization capabilities.
Table 1 shows the scores achieved for each of the settings. While GRU and transformers seem to do a good job at encoding history for the MSD environment, both are significantly worse on the more complex DMSD task when compared to our proposed encoders. This is true especially for GRU, which performs worse than two independent PID controllers for every configuration. Additionally, while it seems that GRU can generalize to large amounts of variation in system parameters when a small amount is present, it fails horribly when trained on fixed system parameters. On the other hand, transformers are able to generalize surprisingly well when trained on both fixed system parameters and with small variation. We hypothesize the autoregressive nature of GRU may make it particularly susceptible to overfitting. Comparing PIDE and GPIDE, we see that PIDE tends to shine in the straightforward cases where there is little change in system parameters, whereas GPIDE is able to adapt when there is a large variation in parameters since it has additional capacity.
Navigation Environment To emulate the setting where the policy is trained on an imperfect simulator, we consider an environment in which the agent is tasked with moving itself across a surface to a specified 2D target as quickly and efficiently as possible. At every point in time, the agent can apply some force to move itself, but a penalty term proportional to the magnitude of the force is subtracted from the reward. Suppose that we have access to a simulator of the environment that is perfect except for the fact that it does not model friction between the agent and the surface. We refer to this simulator and the real environment as the "No Friction" and "Friction" environment, respectively. In both environments, the mass of the agent is treated as a system parameter that is sampled for each episode; however, the Friction environment has a larger range of masses and also randomly samples the coefficient of friction each episode.
Figure 2 shows the average returns recorded during training for both navigation environments and when the policies trained in No Friction are evaluated in Friction. A table of final scores can be found in Appendix G.3. One can see that GPIDE not only achieves the best returns in the environments it was trained in, but is also robust when going from the frictionless environment to the one with friction. On the other hand, PIDE has less capacity and therefore cannot achieve the same results; however, it
Figure 2: Average Returns for Navigation Environments. The curves show the average over five seeds and the shaded region shows the standard error. For this plot, we allowed for  $5\mathrm{x}$  the normal amount budget to allow all methods to converge. We omit the PID controllers from this plot since it gets substantially worse returns.
is immediately more robust than the other methods, although it begins to overfit over time. It is also clear that using GRU is less sample efficient and less robust to changes in the test environment.
Tokamak Control For our last tracking experiment we return to tokamak control. In particular, we focus on the DIII-D tokamak, a device operated by General Atomics in San Diego, California. We aim to control two quantities:  $\beta_{N}$ , the normalized ratio between plasma and magnetic pressure, and rotation, i.e. how fast the plasma is spinning around the toroid. These are important quantities to track because  $\beta_{N}$  serves as an approximate economic indicator and rotation control of the plasma has been suggested to be key for stability [7, 67, 10, 61, 56]. The policy has control over the eight neutral beams [24], which are able to inject power and torque by blasting neutrally charged particles into the plasma. Importantly, two of the eight beams can be oriented in the opposite direction from the others, which decouples the total combined power and torque to some extent (see Figure 3).
To emulate the sim-to-real training experience, we create a simulator based on the equations described in Boyer et al. [9] and Scoville et al. [63]. This simulator has two major shortcomings: it assumes that certain states of the plasma (e.g. its shape) are fixed for entire episodes, and it assumes that there are no events that cause loss of confinement of the plasma. We make up for part of the former by randomly sampling plasma states each episode. The approximate
Figure 3: Illustration of DIII-D from Above. [14] Each beamline in the figure contains two independent beams (yellow boxes). The plasma is rotating counter-clockwise and the two beams in the bottom left of the figure are oriented in the counter-current direction, allowing power and torque to be decoupled. This figure gives a rough idea of beam positioning but is not physically accurate.
Environment (Train/Test)	PID Controller	GRU	Transformer	PIDE	GPIDE
βN-Track Sim/Sim	40.69 ± 0.32	100.00 ± 0.20	97.56 ± 0.19	0.00 ± 1.05	98.33 ± 0.41
βN-Track Sim/Real	89.15 ± 0.99	40.96 ± 5.45	40.05 ± 11.91	0.00 ± 21.04	55.21 ± 4.44
βN-Track Real/Real	98.45 ± 0.77	98.24 ± 0.38	98.74 ± 0.29	100.00 ± 0.23	99.30 ± 0.64
Average	76.10	79.73	78.79	33.33	84.28
βN-Rot-Track Sim/Sim	0.00 ± 0.83	99.06 ± 0.22	96.22 ± 0.94	67.98 ± 0.50	100.00 ± 0.29
βN-Rot-Track Sim/Real	83.71 ± 2.64	39.76 ± 5.84	33.31 ± 0.69	0.00 ± 8.89	51.00 ± 1.92
βN-Rot-Track Real/Real	92.02 ± 0.84	98.34 ± 0.52	96.32 ± 0.31	98.21 ± 0.23	100.00 ± 0.46
Average	58.58	79.05	75.28	55.40	83.67
Total Average	67.34	79.39	77.03	44.36	83.97
Table 2: Tokamak Control Task Results. The scores presented are averaged over five seeds and we show the standard error for each score.
Environment	PPO-GRU	TD3-GRU	VRM	SAC-Transformer	SAC-GPIDE
HalfCheetah-P	27.09 ± 7.85	85.80 ± 5.15	-107.00 ± 1.39	37.00 ± 9.97	82.63 ± 3.46
Hopper-P	49.00 ± 5.22	84.63 ± 8.33	3.53 ± 1.63	59.54 ± 19.64	93.27 ± 13.56
Walker-P	1.67 ± 4.39	29.08 ± 9.67	-3.89 ± 1.25	24.89 ± 14.80	96.61 ± 1.60
Ant-P	39.48 ± 3.74	-36.36 ± 3.35	-36.39 ± 0.17	-10.57 ± 2.34	66.66 ± 2.94
HalfCheetah-V	19.68 ± 11.71	59.03 ± 2.88	-80.49 ± 2.97	-41.31 ± 26.15	20.39 ± 29.60
Hopper-V	13.86 ± 4.80	57.43 ± 8.63	10.08 ± 3.51	0.28 ± 8.49	90.98 ± 4.28
Walker-V	8.12 ± 5.43	-4.63 ± 1.30	-1.80 ± 0.70	-8.21 ± 1.31	36.90 ± 16.59
Ant-V	1.43 ± 3.26	17.03 ± 6.55	-13.41 ± 0.12	0.81 ± 1.31	18.03 ± 5.10
Average	20.04	36.50	-28.67	7.80	63.18
Table 3: PyBullet Task Results. Each score is averaged over four seeds and we report the standard errors. Unlike before, we scale the returns by the returns of an oracle policy (i.e. one which sees position and velocity) and a policy which does not encode any history. For the environment names, "P" and "V" denote only position or only velocity in the observation, respectively.
"real" environment addresses these shortcomings by using a data-driven simulator. This approach to simulating has been shown to be relatively accurate [14, 65, 64, 1], and we use an adapted version of the simulator appearing in Char et al. [14] for our work. This simulator accounts for a greater set of the plasma's state, and the additional information is rich enough that loss of confinement events play a role in the dynamics.
We consider two versions of this task: the first is a SISO task where total power is controlled to achieve  $\beta_{N}$  target, and the second is a MIMO task where total power and torque is controlled to achieve  $\beta_{N}$  and rotation targets. The results for both of these tasks are shown in Table 2. Most of the RL techniques are able to do well if tested in the same environment they were trained in; the exception of this is PIDE, which curiously is unable to perform well in the simulator environment. While no reinforcement learning method matches the robustness of a PID controller, policies trained with GPIE fare significantly better.
5.2 PyBullet Locomotion Tasks
Moving past tracking problems, we evaluate GPIDE on the PyBullet [16] benchmark proposed by Han et al. [28] and adapted in Ni et al. [50]. The benchmark has four robots: halfcheetah, hopper, walker, and ant. For each of these, either the current position information or velocity information is hidden from the agent. Except for GPIDE and transformer encoder, we use all of the performance traces given by Ni et al. [50]. In addition to SAC, they also train using PPO [62], A2C [49], TD3 [23], and VRM [28], a variational method that uses recurrent units to estimate the belief state. We reproduce as much of the training and evaluation procedure as possible, including using the same hyperparameters in the SAC algorithm and giving the history encoders access to actions and rewards. For more information see Appendix C.2. Table 3 shows the performance of GPIDE along with a subset of best performing methods (more results can be found in Appendix G.5). These results make it clear that GPIDE is powerful in arbitrary control tasks besides tracking since the average score achieved across all tasks is a  $73\%$  improvement over TD3-GRU, which we believe is the previous state-of-the-art for this benchmark at the time of this work.
Moreover, GPIDE dominates performance for every robot except HalfCheetah. The only setting where GPIDE achieves significantly worse performance is HalfCheetah-V. This setting seemed to cause stability issues in some seeds, lowering the average score. We believe that these stability issues stem from the attention heads, and we found that removing these heads fixed stability issues and resulted in a competitive average score (see Appendix G.5).
5.3 GPIDE Ablations
To investigate the role of each type of head, we reran all experiments with three variants of GPIDE: one with six exponential smoothing heads (ES), one with five exponential smoothing heads and
MSD	DMSD	Navigation	βN Track	βN-Rot Track	PyBullet
ES	+2.69%	-11.14%	-0.11%	+2.57%	+0.29%	+5.81%
ES + Sum	-8.33%	+5.49%	-1.65%	+4.22%	+0.76%	+11.00%
Attention	-0.36%	-54.95%	-3.91%	-8.85%	-7.55%	-39.44%
Table 4: GPIDE Ablation Percent Difference for Average Scores. All final scores can be found in Appendix G.
one summation head (ES + Sum), and one with six attention heads (see Appendix C.3 for details). We choose these three configurations specifically to better understand the roles that attention and summation play.
Table 4 shows the differences in the average scores for each environment. The first notable takeaway is that having summation is often important in some of the more complex environments. The other takeaway is that much of the heavy lifting is being done by the exponential smoothing. GPIDE fares far worse when only having attention heads, especially in DMSD and the PyBullet environments.
Figure 4: Averaged Attention Schemes for MSD-Small and HalfCheetah-P. Each y-position on the grid corresponds to an amount of history being recorded, and each x-position corresponds to a time point in that history. As such, each of the left-most points are the oldest observation in the history, and the diagonals correspond to the most recent observation. The darker the blue, the greater the weight that is assigned to that time point.
We visualize some of the attention schemes learned by GPIDE for MSD with small variation and HalfCheetah (Figure 4). While the attention scheme learned for MSD could potentially be useful since it recalls information from near the beginning of the episode when the most movement is happening, it appears that the attention scheme for HalfCheetah is simply a poor reproduction of exponential smoothing, making
it redundant and suboptimal. In fact, we found this phenomenon to be true across all attention heads and PyBullet tasks. We believe that the periodicity that appears here is due to the oscillatory nature of the problem and lack of positional encoding (although we found including positional encoding degrades performance).
In this work, we introduced the PIDE and GPIDE history encoders to be used for reinforcement learning in partially observable control tasks. Although both are far simpler than prior methods of encoding, they often result in powerful yet robust controllers. We hope that this work inspires the research community to think about how pre-existing control methods can inform architecture choices.
Limitations There are many different ways a control task may be partially observable, and we do not believe that our proposed methods are solutions to all of them. For example, we do not think GPIDE is necessarily suited for tasks where the agent needs to remember events (e.g. picking up a key to unlock a door).
As with any bias, the PID-inspired biases that we propose in this work come at the cost of flexibility. For the experiments considered in this work, this trade off is beneficial and results in better policies. However, it is unclear whether this trade off is always worth making. It is possible that in higher dimensional environments or environments with more complex dynamics that having more flexibility is preferable to our proposed architecture.
Lastly, some tasks may require the policy to act on images as observations. We are optimistic that PIDE and GPIDE are still useful architectures in this setting, but we speculate that this is contingent on training an image encoder that is well-suited for these architectures, and we leave this research direction for future work.
We would like to thank Conor Igoe for his helpful discussions and advice on this work. We would also like to thank the NeurIPS 2023 reviewers assigned to our paper for their discussion and feedback.
This work was supported in part by US Department of Energy grants under contract numbers DE-SC0021414 and DE-AC02-09CH1146. This work is also supported by the National Science Foundation Graduate Research Fellowship Program under Grant No. DGE1745016 and DGE2140739. Any opinions, findings, and conclusions or recommendations expressed in this material are those of the author(s) and do not necessarily reflect the views of the National Science Foundation.
[71] T Wakatsuki, T Suzuki, N Oyama, and N Hayashi. Ion temperature gradient control using reinforcement learning technique. *Nuclear Fusion*, 61(4):046036, 2021.
[72] ML Walker, DA Humphreys, JA Leuer, JR Ferron, and BG Penaflor. Implementation of model-based multivariable control on diiii-d. GA-A23468, 2000.
[73] Jane X Wang, Zeb Kurth-Nelson, Dhruva Tirumala, Hubert Soyer, Joel Z Leibo, Remi Munos, Charles Blundell, Dharshan Kumaran, and Matt Botvinick. Learning to reinforcement learn. arXiv preprint arXiv:1611.05763, 2016.
[74] Junhang Wang, Jianing Zhu, Chao Zou, Linlin Ou, and Xinyi Yu. Robust adaptive pid control based on reinforcement learning for mimo nonlinear six-joint manipulator. In 2021 China Automation Congress (CAC), pages 2633-2638. IEEE, 2021.
[75] Zhihan Yang and Hai Huu Nguyen. Recurrent off-policy baselines for memory-based continuous control. In Deep RL Workshop NeurIPS 2021, 2021.
[76] Wenhao Yu, Jie Tan, C Karen Liu, and Greg Turk. Preparing for the unknown: Learning a universal policy with online system identification. arXiv preprint arXiv:1702.02453, 2017.
[77] Luisa Zintgraf, Kyriacos Shiarlis, Maximilian Igl, Sebastian Schulze, Yarin Gal, Katja Hofmann, and Shimon Whiteson. Varibad: A very good method for bayes-adaptive deep rl via meta-learning. arXiv preprint arXiv:1910.08348, 2019.
1 Introduction 1
2 Preliminaries 2
3 Methodology 3
4 Related Work 5
5.1 Tracking Problems 6
5.2 PyBullet Locomotion Tasks 9
5.3 GPIDE Ablations 9
6 Discussion 10
7 Acknowledgements 10
Appendices 16
A Additional GPIDE Details 17
A.1 Forming the P, I, and D Features with GPIDE 17
B Implementation Details 17
B.1 GPIDE Implementation Details 18
B.2 Recurrent and Transformer Baseline Details 18
B.3 PID Baseline 18
C Hyperparameters 19
C.1 Hyperparameters for Tracking Tasks 19
C.2 Hyperparameters for PyBullet Task 19
C.3 Hyperparameters for Ablations 19
D Computation Details 20
E Environment Descriptions 20
E.1 Mass Spring Damper 20
E.2 Navigation Environment 21
E.3 Tokamak Control Environment 21
F Additional Experiments 23
F.1 Experiments Using VIB + GRU 23
F.2 Lookback Size Ablations 24
G Further Results 24
G.1 Evaluation Procedure 25
G.2 MSD and DMSD Results 25
G.3 Navigation Results 28
G.4 Tokamak Control Results 30
G.5 PyBullet Results 33
G.6 Attention Scheme Visualizations 36
A Additional GPIDE Details
A.1 Forming the P, I, and D Features with GPIDE
As mentioned in Section 3, the P, I, and D features associated with a PID controller can be easily formed with GPIDE, and in this Appendix we show this concretely. Assume that the observations take the form  $o_{t} = \left(x_{t}^{(1)},\dots ,x_{t}^{(M)},\sigma_{t}^{(1)},\dots ,\sigma_{t}^{(M)}\right)$ , and assume that the actions and rewards are not given to GPIDE for simplicity. Then the linear projection for each head,  $f_{\theta}^{h}$ , takes in  $x_{t - 1}^{(1)},\ldots ,x_{t - 1}^{(M)},\sigma_{t - 1}^{(1)},\ldots ,\sigma_{t - 1}^{(M)},(x_{t}^{(1)} - x_{t - 1}^{(1)}),\ldots ,(x_{t}^{(M)} - x_{t - 1}^{(M)}),(\sigma_{t}^{(1)} - \sigma_{t - 1}^{(1)}),\ldots ,(\sigma_{t}^{(M)} - \sigma_{t - 1}^{(M)})$  at each time step  $t$ . We use a GPIDE architecture with three heads and where  $g_{\theta}$  is the identity function. The linear projections for each head is as follows:
$$
\begin{array}{l} f _ {\theta} ^ {1} (o _ {t - 1}, o _ {t} - o _ {t - 1}) = f _ {\theta} ^ {2} (o _ {t - 1}, o _ {t} - o _ {t - 1}) = \left[ \begin{array}{c} (x _ {t} ^ {(1)} - x _ {t - 1} ^ {(1)}) + x _ {t - 1} ^ {(1)} - (\sigma_ {t} ^ {(1)} - \sigma_ {t - 1} ^ {(1)}) + \sigma_ {t - 1} ^ {(1)} \\ \vdots \\ (x _ {t} ^ {(M)} - x _ {t - 1} ^ {(M)}) + x _ {t - 1} ^ {(M)} - (\sigma_ {t} ^ {(M)} - \sigma_ {t - 1} ^ {(M)}) + \sigma_ {t - 1} ^ {(M)} \end{array} \right] \\ f _ {\theta} ^ {3} (o _ {t - 1}, o _ {t} - o _ {t - 1}) = \left[ \begin{array}{c} (x _ {t} ^ {(1)} - x _ {t - 1} ^ {(1)}) - (\sigma_ {t} ^ {(1)} - \sigma_ {t - 1} ^ {(1)}) \\ \vdots \\ (x _ {t} ^ {(M)} - x _ {t - 1} ^ {(M)}) - (\sigma_ {t} ^ {(M)} - \sigma_ {t - 1} ^ {(M)}) \end{array} \right] \\ \end{array}
$$