SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Figure 9: Removing terms from the value function input. Learning curves on the MuJoCo benchmark, when removing components from the input to the value function. Results are averaged over 10 seeds. The shaded area captures a  $95\%$  confidence interval around the average performance.
Figure 10: Exclusive value function input. Learning curves on the MuJoCo benchmark, changing the value function input to a single component. Results are averaged over 10 seeds. The shaded area captures a  $95\%$  confidence interval around the average performance.
Figure 11: Modified policy input. Learning curves on the MuJoCo benchmark, changing the policy input to a single component. Results are averaged over 10 seeds. The shaded area captures a  $95\%$  confidence interval around the average performance.
Figure 12: Fixed embeddings. Learning curves on the MuJoCo benchmark, without fixed embeddings for the value function and policy, so that the current encoder is used rather than the fixed encoder. Results are averaged over 10 seeds. The shaded area captures a  $95\%$  confidence interval around the average performance.
D.3 Normalization
SALE normalizes embeddings and input through our proposed AvgL1Norm:
$$
\operatorname {A v g L 1 N o r m} (x) := \frac {x}{\frac {1}{N} \sum_ {i} | x _ {i} |}. \tag {46}
$$
AvgL1Norm is embedded into the architecture used by SALE (see Pseudocode 1), applied to the output of the state encoder  $f$  to give the state embedding  $z^{s} = \mathrm{AvgL1Norm}(f(s))$ . It is also used following a linear layer in the input of both the value function and policy. This makes the input to the value function and policy:
$$
Q \left(z ^ {s a}, \operatorname {A v g L 1 N o r m} (f (s)), \operatorname {A v g L 1 N o r m} \left(\phi^ {s a}\right)\right), \tag {47}
$$
$$
\pi (\operatorname {A v g L 1 N o r m} (f (s)), \operatorname {A v g L 1 N o r m} (\phi^ {s})), \tag {48}
$$
where  $\phi^{sa}$  is the output from a linear layer on the state-action input,  $\operatorname{Linear}(s,a)$ , and  $\phi^s$  is the output from a linear layer on the state input,  $\operatorname{Linear}(s)$ .
In this section, we vary how normalization is used, swapping which components it is used on and what normalization function is used.
(Figure 13) Normalization 1: No normalization on  $\phi$ . Normalization is not applied to either  $\phi$ . This makes the input to the value function and policy:
$$
Q \left(z ^ {s a}, \operatorname {A v g L 1 N o r m} (f (s)), \phi^ {s a}\right), \tag {49}
$$
$$
\pi (\operatorname {A v g L 1 N o r m} (f (s)), \phi^ {s}). \tag {50}
$$
(Figure 13) Normalization 2: No normalization. Normalization is not applied anywhere. This makes the input to the value function and policy:
$$
Q \left(z ^ {s a}, f (s), \phi^ {s a}\right), \tag {51}
$$
$$
\pi (f (s), \phi^ {s}). \tag {52}
$$
(Figure 14) Normalization 3: Normalization used on  $z^{sa}$ . Since  $z^{sa}$  is trained to predict the next state embedding  $z^{s'}$ , which is normalized, normalization may not be needed for the output of the state-action encoder  $g$ . We try also applying normalization to the output of  $g$ . This makes the input to the value function and policy:
$$
Q \left(\operatorname {A v g L 1 N o r m} \left(g \left(z ^ {s}, a\right)\right), \operatorname {A v g L 1 N o r m} (f (s)), \operatorname {A v g L 1 N o r m} \left(\phi^ {s a}\right)\right), \tag {53}
$$
$$
\pi (\operatorname {A v g L 1 N o r m} (f (s)), \operatorname {A v g L 1 N o r m} (\phi^ {s})). \tag {54}
$$
(Figure 15) Normalization 4: BatchNorm. BatchNorm is used instead of AvgL1Norm. This makes the input to the value function and policy:
$$
Q \left(z ^ {s a}, \operatorname {B a t c h N o r m} (f (s)), \operatorname {B a t c h N o r m} \left(\phi^ {s a}\right)\right), \tag {55}
$$
$$
\pi (\operatorname {B a t c h N o r m} (f (s)), \operatorname {B a t c h N o r m} \left(\phi^ {s}\right)), \tag {56}
$$
(Figure 15) Normalization 5: LayerNorm. LayerNorm is used instead of AvgL1Norm. This makes the input to the value function and policy:
$$
Q \left(z ^ {s a}, \text {L a y e r N o r m} (f (s)), \text {L a y e r N o r m} \left(\phi^ {s a}\right)\right), \tag {57}
$$
$$
\pi \left(\operatorname {L a y e r N o r m} (f (s)), \operatorname {L a y e r N o r m} \left(\phi^ {s}\right)\right), \tag {58}
$$
(Figure 16) Normalization 3: Cosine similarity loss. This is inspired by SPR [Schwarzer et al., 2020], where normalization is not used and a cosine similarity loss is used instead of MSE for updating the encoder. This means the encoder loss function, which is originally:
$$
\mathcal {L} (f, g) := \left(g (f (s), a) - | f (s ^ {\prime}) | _ {\times}\right) ^ {2} = \left(z ^ {s a} - | z ^ {s ^ {\prime}} | _ {\times}\right) ^ {2}, \tag {59}
$$
now becomes
$$
\mathcal {L} (f, g) := \operatorname {C o s i n e} \left(z ^ {s a}, \left| z ^ {s ^ {\prime}} \right| _ {\times}\right), \tag {60}
$$
$$
\operatorname {C o s i n e} \left(z ^ {s a}, \left| z ^ {s ^ {\prime}} \right| _ {\times}\right) := \left(\frac {z ^ {s a}}{\| z ^ {s a} \| _ {2}}\right) ^ {\top} \left(\frac {\left| z ^ {s ^ {\prime}} \right| _ {\times}}{\| \left| z ^ {s ^ {\prime}} \right| _ {\times} \| _ {2}}\right). \tag {61}
$$
No normalization is used elsewhere.
Table 10: Average performance on the MuJoCo benchmark at 1M time steps.  $\pm$  captures a  $95\%$  confidence interval around the average performance. Results are over 10 seeds.
Algorithm	HalfCheetah	Hopper	Walker2d	Ant	Humanoid
TD7 (no checkpoints)	17123 ± 296	3361 ± 429	5718 ± 308	8605 ± 1008	7381 ± 172
TD3	10574 ± 897	3226 ± 315	3946 ± 292	3942 ± 1030	5165 ± 145
No normalization on φ	17231 ± 246	2647 ± 466	5639 ± 1248	8191 ± 846	5862 ± 1471
No normalization	17275 ± 288	3359 ± 479	6168 ± 164	7274 ± 662	4803 ± 1706
Normalization on zsa	16947 ± 284	3383 ± 262	5502 ± 1142	8049 ± 538	6418 ± 302
BatchNorm	17299 ± 218	2318 ± 473	4839 ± 987	7636 ± 769	6979 ± 325
LayerNorm	17132 ± 360	2451 ± 902	4470 ± 1096	6331 ± 1024	6712 ± 1297
Cosine similarity loss	16897 ± 372	3324 ± 249	5566 ± 353	7873 ± 511	4370 ± 1573
Figure 13: Removing normalization. Learning curves on the MuJoCo benchmark, where normalization is removed on certain components of the input to the value function and policy. Results are averaged over 10 seeds. The shaded area captures a  $95\%$  confidence interval around the average performance.
Figure 14: Normalization on  $z^{sa}$ . Learning curves on the MuJoCo benchmark, where normalization is also applied to the state-action embedding  $z^{sa}$ . Results are averaged over 10 seeds. The shaded area captures a  $95\%$  confidence interval around the average performance.
Figure 15: Alternate normalization. Learning curves on the MuJoCo benchmark, where BatchNorm or LayerNorm is used instead of AvgL1Norm. Results are averaged over 10 seeds. The shaded area captures a  $95\%$  confidence interval around the average performance.
Figure 16: Cosine similarity instead of normalization. Learning curves on the MuJoCo benchmark, where normalization is removed and replaced with a cosine similarity loss, as inspired by SPR [Schwarzer et al., 2020]. Results are averaged over 10 seeds. The shaded area captures a  $95\%$  confidence interval around the average performance.
D.4 End-to-end
The embeddings used in SALE are decoupled, meaning that the encoders used to output the embeddings are trained independently from the value function or policy. Instead, we might consider training the encoders end-to-end with the value function. Doing so requires training the current embedding with the value function loss (rather than using the fixed embedding). This means that the encoders are trained with the following loss:
$$
\mathcal {L} \left(f _ {t + 1}, g _ {t + 1}\right) := \operatorname {H u b e r} \left(\operatorname {t a r g e t} - Q _ {t + 1} \left(z _ {t + 1} ^ {s a}, z _ {t + 1} ^ {s}, s, a\right)\right) + \beta \left(z _ {t + 1} ^ {s a} - \left| z _ {t + 1} ^ {s ^ {\prime}} \right| _ {\times}\right) ^ {2}, \tag {62}
$$
$$
\text {t a r g e t} := r + \gamma \operatorname {c l i p} \left(\min  \left(Q _ {t, 1} (x), Q _ {t, 2} (x)\right), Q _ {\min }, Q _ {\max }\right), \tag {63}
$$
$$
x := \left[ z _ {t} ^ {s ^ {\prime} a ^ {\prime}}, z _ {t} ^ {s ^ {\prime}}, s ^ {\prime}, a ^ {\prime} \right], \tag {64}
$$
$$
a ^ {\prime} := \pi_ {t} \left(z _ {t} ^ {s ^ {\prime}}, s ^ {\prime}\right) + \epsilon , \tag {65}
$$
$$
\epsilon \sim \operatorname {c l i p} \left(\mathcal {N} (0, \sigma^ {2}), - c, c\right). \tag {66}
$$
In Table 11 and Figure 17 we display the performance when varying the hyperparameter  $\beta$ .
Table 11: Average performance on the MuJoCo benchmark at 1M time steps.  $\pm$  captures a  $95\%$  confidence interval around the average performance. Results are over 10 seeds.
Algorithm	HalfCheetah	Hopper	Walker2d	Ant	Humanoid
TD7 (no checkpoints)	17123 ± 296	3361 ± 429	5718 ± 308	8605 ± 1008	7381 ± 172
TD3	10574 ± 897	3226 ± 315	3946 ± 292	3942 ± 1030	5165 ± 145
End-to-end, 0.1	16186 ± 360	1820 ± 674	5013 ± 729	6601 ± 1251	5076 ± 897
End-to-end, 1	15775 ± 658	1779 ± 537	4882 ± 710	6604 ± 1135	5880 ± 703
End-to-end, 10	16472 ± 365	1534 ± 341	4900 ± 760	6626 ± 1253	5279 ± 1095
Figure 17: End-to-end learning. Learning curves on the MuJoCo benchmark, when training the embeddings end-to-end with the value function, where the encoder loss is weighted with respect to  $\beta$ . Results are averaged over 10 seeds. The shaded area captures a  $95\%$  confidence interval around the average performance.
E Extrapolation Error Study
E.1 Extrapolation Error
Extrapolation error is an issue with deep RL where the value function poorly estimates the value of unseen or rare actions [Fujimoto et al., 2019]. If we assume that  $(s, a)$  is not contained in the dataset, then  $Q(s, a)$  is effectively a guess generated by the value function  $Q$ . During training,  $s$  is always sampled from the dataset, but  $a$  is often generated by the policy  $(a \sim \pi(s))$ , resulting in a potentially poor estimate.
Extrapolation error is typically considered a problem for offline RL, where the agent is given a fixed dataset and cannot interact further with the environment, as actions sampled from the policy may not be contained in the dataset. Extrapolation error is not considered a problem in online RL since the policy interacts with the environment, collecting data for the corresponding actions it generates.
In our empirical analysis (Section E.2) we observe the presence of extrapolation error when using SALE. Our hypothesis is that by significantly expanding the input dimension dependent on the action, the network becomes more prone to erroneous extrapolations (for example, for the Ant environment, the original action dimension size is 8, but the state-action embedding  $z^{sa}$  has a dimension size of 256). This is because unseen actions can appear significantly more distinct from seen actions, due to the dramatic increase in the number of dimensions used to represent the action-dependent input.
E.2 Empirical Analysis
In this section we vary the input to the value function in TD7 to understand the role of the input dimension size and extrapolation error. The default input to the value function in TD7 is as follows:
$$
Q \left(z ^ {s a}, z ^ {s}, \phi\right), \tag {67}
$$
$$
\phi := \operatorname {A v g L 1 N o r m} (\operatorname {L i n e a r} (s, a)). \tag {68}
$$
Throughout all experiments, we assume no value clipping (as this is introduced in response to the analysis from this section). This makes the value function loss (originally Equation 15):
$$
\mathcal {L} \left(Q _ {t + 1}\right) := \operatorname {H u b e r} \left(\operatorname {t a r g e t} - Q _ {t + 1} \left(z _ {t} ^ {s a}, z _ {t} ^ {s}, s, a\right)\right), \tag {69}
$$
$$
\text {t a r g e t} := r + \gamma \min  \left(Q _ {t, 1} (x), Q _ {t, 2} (x)\right), \tag {70}
$$
$$
x := \left[ z _ {t - 1} ^ {s ^ {\prime} a ^ {\prime}}, z _ {t - 1} ^ {s ^ {\prime}}, s ^ {\prime}, a ^ {\prime} \right], \tag {71}
$$
$$
a ^ {\prime} := \pi_ {t} \left(z _ {t - 1} ^ {s ^ {\prime}}, s ^ {\prime}\right) + \epsilon , \tag {72}
$$
$$
\epsilon \sim \operatorname {c l i p} \left(\mathcal {N} \left(0, \sigma^ {2}\right), - c, c\right). \tag {73}
$$
Policy checkpoints are not used and the policy is not modified. Since extrapolation error is closely linked to available data, we also vary the maximum size of the replay buffer (default: 1M).
(Figure 18) Extrapolation 1: No clipping. This is TD7 without value clipping in the value function loss as discussed above.
(Figure 19) Extrapolation 2: No  $z^{sa}$ . We remove  $z^{sa}$  from the value function input. This makes the input:  $Q_{t + 1}(z_t^s, \phi)$ .
(Figure 20) Extrapolation 3: Small  $\phi$ . We reduce the number of dimensions of  $\phi$  from 256 to 16.
(Figure 21) Extrapolation 4: No  $z^{sa}$  and small  $\phi$ . We remove  $z^{sa}$  from the value function input and reduce the number of dimensions of  $\phi$  from 256 to 16.
(Figure 22) Extrapolation 5: Frozen embeddings. The encoders are left unchanged throughout training, by leaving them untrained. This means the input of the value function is  $Q_{t + 1}(z_0^s,\phi)$ .
(Figure 23) Extrapolation 6: TD7. The full TD7 method (without checkpoints).
Figure 18: No clipping. The performance curve and corresponding value estimate made by the value function, where no value clipping is used. The results for each individual seed are presented for a replay buffer size of either 1M (left) or 100k (right). The environment used is Ant.
Figure 19: No  $z^{sa}$ . The performance curve and corresponding value estimate made by the value function, where  $z^{sa}$  is not included in the input. The results for each individual seed are presented for a replay buffer size of either 1M (left) or 100k (right). The environment used is Ant.
Figure 20: Small  $\phi$ . The performance curve and corresponding value estimate made by the value function, where the dimension size of  $\phi$  is reduced. The results for each individual seed are presented for a replay buffer size of either 1M (left) or 100k (right). The environment used is Ant.
Figure 21: No  $z^{sa}$ , small  $\phi$ . The performance curve and corresponding value estimate made by the value function, where  $z^{sa}$  is not included in the input and the dimension size of  $\phi$  is reduced. The results for each individual seed are presented for a replay buffer size of either 1M (left) or 100k (right). The environment used is Ant.
Figure 22: Frozen embeddings. The performance curve and corresponding value estimate made by the value function, where the encoders are not trained. The results for each individual seed are presented for a replay buffer size of either 1M (left) or 100k (right). The environment used is Ant.
Figure 23: TD7. The performance curve and corresponding value estimate made by the value function of TD7 (using value clipping). The results for each individual seed are presented for a replay buffer size of either 1M (left) or 100k (right). The environment used is Ant.
F Policy Checkpoints
F.1 Motivation
Deep RL algorithms are widely known for their inherent instability, which often results in substantial variance in performance during training. Instability in deep RL can occur on a micro timescale (performance can shift dramatically between episodes) [Henderson et al., 2017, Fujimoto and Gu, 2021] and a macro timescale (the algorithm can diverge or collapse with too much training) [Kumar et al., 2020a, Lyle et al., 2021]. In Figure 24 & 25 we show the instability of RL methods by (a) showing how presenting the average performance can hide the instability in a single trial of RL (Figure 24), and (b) measuring how much the performance can drop between nearby evaluations (Figure 25).
Figure 24: Standard presentation Single seed
The training curve as commonly presented in RL papers and the training curve of a single seed. Both curves are from the TD3 algorithm, trained for 5M time steps. The standard presentation is to evaluate every  $N_{\mathrm{freq}}$  steps, average scores over  $N_{\mathrm{episodes}}$  evaluation episodes and  $N_{\mathrm{seeds}}$  seeds, then to smooth the curve by averaging over a window of  $N_{\mathrm{window}}$  evaluations. (In our case this corresponds to  $N_{\mathrm{freq}} = 5000$ ,  $N_{\mathrm{episodes}} = 10$ ,  $N_{\mathrm{seeds}} = 10$ ,  $N_{\mathrm{window}} = 10$ ). The learning curve of a single seed has no smoothing over seeds or evaluations ( $N_{\mathrm{seeds}} = 1$ ,  $N_{\mathrm{window}} = 1$ ). By averaging over many seeds and evaluations, the training curves in RL can appear deceptively smooth and stable.
Figure 25: window = 1 window = 5 window = 10
The  $\%$  of evaluations which suffer from a performance drop of at least  $x\%$  in the next window evaluations for the TD3 algorithm, trained for 5M time steps. For example, for any given evaluation, the likelihood of one of the next 10 evaluations performing at least  $50\%$  worse is around  $30\%$ . Only evaluations from the last 1M training steps are used, so that early training instability is not accounted for. This shows that large drops in performance are common.
F.2 Method
The overview for our approach for using checkpoints in RL:
- Assess the current policy (using training episodes).
- Train the current policy (with a number of time steps equal (or proportional) to the number of time steps viewed during assessment).
- If the current policy outperforms the checkpoint policy, then update the checkpoint policy.
Note that we aim to use training episodes to evaluate the checkpoint policy, as any evaluation or test time episodes are considered entirely separate from the training process. For clarity, we will use the term assessment (rather than evaluation) to denote any measure of performance which occurs during training. In Algorithm 3 we outline the basic policy checkpoint strategy.
Minimum performance criteria. An interesting choice that remains is how to determine whether the current policy outperforms the checkpoint policy. Instead of using the average performance, we consider the minimum performance over a small number of assessment episodes. Using the minimum provides a criteria which is more sensitive to instability, favoring policies which achieve a consistent performance.
Additionally, using the minimum allows us to end the assessment phase early, if the minimum performance of the current policy drops below the minimum performance of the checkpoint policy. This allows us to use a higher maximum number of assessment episodes, while not wasting valuable training episodes assessing suboptimal policies. In Algorithm 4 we outline policy checkpoints where the minimum performance is used, and the assessment phase is variable length due to early termination.
Algorithm 3 Basic Policy Checkpoints
1: for episode = 1 to assessment Episodes do ▷ Assessment  
2: Follow the current policy  $\pi_{t+1}$  and determine episode reward.  
3: Update the performance measure of the current policy.  
4: Increment timesteps_since_training by the length of the episode.  
5: if actor  $\pi_{t+1}$  outperforms checkpoint policy  $\pi_c$  then ▷ Checkpointing  
6: Update checkpoint networks  $\pi_c \gets \pi_{t+1}, f_c \gets f_t$ .  
7: Update checkpoint performance.  
8: for  $i = 1$  to timesteps_since_training do ▷ Training  
9: Train RL agent.  
10: Reset the performance measure of the current policy.
Algorithm 4 Policy Checkpoints with Minimum Performance and Early Termination
1: for episode = 1 to assessment Episodes do ▷ Assessment  
2: Follow the current policy  $\pi_{t+1}$  and determine episode reward.  
3: min_performance ← min(min_performance, episode Reward).  
4: Increment timesteps_since_training by the length of the episode.  
5: if min_performance ≤ checkpoint_performance then ▷ Early termination  
6: End current assessment.  
7: if min_performance ≥ checkpoint_performance then ▷ Checkpointing  
8: Update checkpoint networks  $\pi_c \gets \pi_{t+1}, f_c \gets f_t$ .  
9: checkpoint_performance ← min_performance  
10: for i = 1 to timesteps_since_training do ▷ Training  
11: Train RL agent.  
12: Reset min_performance.
Early learning. In our empirical analysis we observe that long assessment phases had a negative impact on early learning (Table 13). A final adjustment to policy checkpoints is to keep the number of assessment episodes fixed to 1 during the early stages of learning, and then increase it to a much larger value later in learning. Since the policy or environment may be stochastic, when increasing the number of assessment episodes, we reduce the performance of the checkpoint policy (by a factor of 0.9), since the checkpoint performance may be overfit to the single episode used to assess its performance.
In Algorithm 1 in the main text, several variables are referenced. These are as follows:
- checkpoint_condition: The checkpoint condition refers to either (a) the maximum number of assessment episodes being reached (20) or the current minimum performance dropping below the minimum performance of the checkpoint policy.
- outperforms: The maximum number of assessment episodes were reached and the current minimum performance is higher than the minimum performance of the checkpoint policy.
Variable descriptions. In Table 3 (in Section A.2), which describes the hyperparameters of TD7, several variables are referenced. These are as follows:
- Checkpoint criteria: the measure used to evaluate the performance of the current and checkpoint policy (default: minimum performance over the assessment episodes).
- Early assessment episodes: the maximum number of assessment episodes during the early stage of learning (default: 1).
- Late assessment episodes: the maximum number of assessment episodes during the late stage of learning (default: 20).
- Early time steps: the duration of the early learning stage in time steps. We consider early learning to be the first 750k time steps.
- Criteria reset weight: The multiplier on the current performance of the checkpoint policy that is applied once, after the early learning stage ends (default: 0.9).
F.3 Empirical Analysis
TD7 uses a checkpointing scheme which:
- Measures the performance of the current and checkpoint policy by the minimum performance over 20 assessment episodes.
- Uses a variable length assessment phase, where the assessment phase terminates early if the minimum performance of the current policy is less than the minimum performance of the checkpoint policy.
- Uses an early learning stage where the maximum number of assessment episodes starts at 1 and is increased to 20 after 750k time steps.
In this section, we vary some of the choices made when using policy checkpoints and analyze the stability benefits of using checkpoints.
(Figure 26) **Checkpoint 1:** Maximum number of assessment episodes. We vary the maximum number episodes used to assess the performance of the current policy (default 20). Early termination of the assessment phase is still used.
(Figure 27) **Checkpoint 2:** Mean, fixed assessment length. We use the mean performance instead of the minimum (default) to determine if the current policy outperforms the checkpoint policy. In this case, early termination of the assessment phase is not used. The number of assessment episodes is fixed at 20.
(Figure 27) **Checkpoint 3:** Mean, variable assessment length. The mean performance is used to determine if the current policy outperforms the checkpoint policy. Early termination of the assessment phase is used. After each episode, we terminate the assessment phase if the current mean performance is below the checkpoint performance. This means the number of assessment episodes is variable, up to a maximum of 20.
(Figure 28) Checkpoint 4: Fixed assessment length. The same number of assessment episodes is used during each assessment phase, without early termination.
(Figure 29) **Checkpoint 5:** Immediate. The number of assessment episodes starts at the corresponding number (as opposed to starting at 1 and increasing to a higher number after the early learning state completes at 750k time steps).