SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Pseudocode 1. TD7 Network Details
Variables:  
zs_dim = 256  
Value Q Network:  
 $\triangleright$  TD7 uses two value networks each with the same network and forward pass.  
10 = Linear(state_dim + action_dim, 256)  
11 = Linear(zs_dim * 2 + 256, 256)  
12 = Linear(256, 256)  
13 = Linear(256, 1)  
Value Q Forward Pass:  
input = concatenate([state, action])  
x = AvgL1Norm(10(input))  
x = concatenate([zsa, zs, x])  
x = ELU(11(x))  
x = ELU(12(x))  
value = 13(x)  
Policy  $\pi$  Network:  
10 = Linear(state_dim, 256)  
11 = Linear(zs_dim + 256, 256)  
12 = Linear(256, 256)  
13 = Linear(256, action_dim)  
Policy  $\pi$  Forward Pass:  
input = state  
x = AvgL1Norm(10(input))  
x = concatenate([zs, x])  
x = ReLU(11(x))  
x = ReLU(12(x))  
action = tanh(13(x))  
State Encoder f Network:  
11 = Linear(state_dim, 256)  
12 = Linear(256, 256)  
13 = Linear(256, zs_dim)  
State Encoder f Forward Pass:  
input = state  
x = ELU(11(input))  
x = ELU(12(x))  
zs = AvgL1Norm(13(x))  
State-Action Encoder g Network:  
11 = Linear(action_dim + zs_dim, 256)  
12 = Linear(256, 256)  
13 = Linear(256, zs_dim)  
State-Action Encoder g Forward Pass:  
input = concatenate([action, zs])  
x = ELU(11(input))  
x = ELU(12(x))  
zsa = 13(x)
Environment. Our experimental evaluation is based on the MuJoCo simulator [Todorov et al., 2012] with tasks defined by OpenAI gym [Brockman et al., 2016] using the v4 environments. No modification are made to the state, action, or reward space.
Terminal transitions. All methods use a discount factor  $\gamma = 0.99$  for non-terminal transitions and  $\gamma = 0$  for terminal transitions. The final transition from an episode which ends due to a time limit is not considered terminal.
Exploration. To fill the replay buffer before training, all methods initially collect data by following a uniformly random policy for the first 25k time steps (256 for TQC). This initial data collection is accounted for in all graphs and tables, meaning "time steps" refers to the number of environment interactions, rather than the number of training steps. Methods based on the deterministic TD3 add Gaussian noise to the policy. Methods based on the stochastic SAC do not add noise.
Evaluation. Agents are evaluated every 5000 time steps, taking the average undiscounted sum of rewards over 10 episodes. Each experiment is repeated over 10 seeds. Evaluations use a deterministic policy, meaning there is no noise added to the policy and stochastic methods use the mean action. Methods with our proposed policy checkpoints use the checkpoint policy during evaluations. Evaluations are considered entirely independent from training, meaning no data is saved, nor is any network updated.
Visualization and tables. As aforementioned, time steps in figures and tables refers to the number of environment interactions (all methods are trained once per environment interaction, other than the initial random data collection phase). The value reported by the table corresponds to the average evaluation over the 10 seeds, at the given time step (300k, 1M or 5M for online results and 1M for offline results). The shaded area in figures and the  $\pm$  term in tables refers to a  $97.5\%$  confidence interval. Given that methods are evaluated over 10 seeds, this confidence interval CI is computed by
$$
\mathrm {C I} = \frac {1 . 9 6}{\sqrt {1 0}} \sigma , \tag {32}
$$
where  $\sigma$  is the sample standard deviation with Bessel's correction at the corresponding evaluation. Unless stated otherwise, all curves are smoothed uniformly over a window of 10 evaluations.
Offline Experiments. Offline results are based on the D4RL datasets [Fu et al., 2021], using the v2 version. No modification is made to the state, action, or reward space. The reported performance is from a final evaluation occurring after 1M training steps, which uses the average D4RL score over 10 episodes. TD7 results are repeated over 10 seeds. Baseline results are taken from other papers, and may not be based on an identical evaluation protocol in terms of number of episodes and seeds used.
Software. We use the following software versions:
- Python 3.9.13
Pytorch 2.0.0 [Paszke et al., 2019]
- CUDA version 11.8
Gym 0.25.0 [Brockman et al., 2016]
- MuJoCo 2.3.3 [Todorov et al., 2012]
C.1 TD3 & TD3+ OFE Hyperparameters
Our TD3 baseline [Fujimoto et al., 2018] uses the author implementation https://github.com/sfujim/TD3. Our TD3+OFE baseline [Ota et al., 2020] uses the aforementioned TD3 code alongside a re-implementation of OFENet. For hyperparameters, the action space is assumed to be in the range  $[-1,1]$ .
Environment-specific hyperparameters. OFE uses a variable number of encoder layers depending on the environment. However, the authors also provide an offline approach for hyperparameter tuning based on the representation loss. Our choice of layers per environment is the outcome of their hyperparameter tuning. Additionally, when training the state-action encoder, OFE drops certain state dimensions which correspond to external forces which are hard to predict. In practice this only affects Humanoid, which uses 292 out of a possible state 376 dimensions on Humanoid. All other tasks use the full state space.
Hyperparameters of both methods are listed in Table 4. Network architecture details of TD3 are described in Pseudocode 2. Network architecture details of TD3+OFE are described in Pseudocode 3 and Pseudocode 4.
Table 4: TD3 & TD3+OFE Hyperparameters.
Hyperparameter	Value
OFE [Ota et al., 2020]	Encoder layers
HalfCheetah	8
Hopper	6
Walker2d	6
Ant	6
Humanoid	8
TD3 [Fujimoto et al., 2018]	Target policy noise σ	N(0, 0.22)
Target policy noise clipping c	(-0.5, 0.5)
Policy update frequency	2
Exploration	Initial random exploration time steps	25k
Exploration noise	N(0, 0.12)
Common	Discount factor γ	0.99
Replay buffer capacity	1M
Mini-batch size	256
Target update rate τ	0.005
Optimizer	(Shared) Optimizer	Adam [Kingma and Ba, 2014]
(Shared) Learning rate	3e-4
Pseudocode 2. TD3 Network Details
Value  $Q$  Network:  $\triangleright$  TD3 uses two value networks each with the same network and forward pass.
l1 = Linear(state_dim + action_dim, 256)  
l2 = Linear(256, 256)  
l3 = Linear(256, 1)
Value  $Q$  Forward Pass:
input  $=$  concatenate(state,action)]   
 $\mathbf{x} =$  ReLU(l1(input))   
 $\mathbf{x} =$  ReLU(12(x))   
value  $= 13(x)$
Policy  $\pi$  Network:
Policy  $\pi$  Forward Pass:
l1 = Linear(state_dim, 256)  
l2 = Linear(256, 256)  
l3 = Linear(256, action_dim)
input  $=$  state   
x  $\equiv$  ReLU(11(input))   
 $\mathbf{x} = \mathbf{Re}\mathbf{L}\mathbf{U}(12(\mathbf{x}))$    
action  $=$  tanh(13(x)
Pseudocode 3. TD3+OFE Network Details (TD3)
Variables:
Value  $Q$  Network:
$\mathtt{zs\_dim} = 240$    
zsa_dim  $= 240$    
num_layers  $= 6$  (or 8)   
 $\mathtt{zs\_hdim} = \mathtt{zs\_dim / num\_layers} = 40$  (or 30)   
zsa_hdim  $=$  zsa_dim/num_layers  $= 40$  (or 30)   
target_dim  $=$  state_dim if not Humanoid else 292
$\triangleright$  TD3+OFE uses two value networks each with the same network and forward pass.
Value  $Q$  Forward Pass:
l1 = Linear(state_dim + action_dim + zs_dim + zsa_dim, 256)  
l2 = Linear(256, 256)  
l3 = Linear(256, 1)
input = zsa  
x = ReLU(11(input))  
x = ReLU(12(x))  
value = 13(x)
Policy  $\pi$  Network:
Policy  $\pi$  Forward Pass:
l1 = Linear(state_dim + zs_dim, 256)  
l2 = Linear(256, 256)  
l3 = Linear(256, action_dim)
input = zs  
x = ReLU(11(input))  
x = ReLU(12(x))  
action = tanh(13(x))
Pseudocode 4. TD3+OFE Network Details (OFE)
State Encoder  $f^s$  Network:   
11 = Linear(state_dim, zs_hdim)   
12 = Linear(state_dim + zs_hdim,zs_hdim)   
13 = Linear(state_dim + zs_hdim * 2,zs_hdim)   
...   
l_num_layers = Linear(state_dim + zs_hdim * (num_layers-1),zs_hdim)   
State Encoder  $f^s$  Forward Pass:   
input = state   
x = swish(batchnorm(11(input)))   
input = concatenate([input,x])   
x = swish(batchnorm(12(input)))   
...   
x = swish(batchnorm(1_num_layers(input)))   
output = concatenate([input,x])   
State-Action Encoder  $f^{sa}$  Network:   
11 = Linear(state_dim + action_dim + zs_dim, zsa_hdim)   
12 = Linear(state_dim + action_dim + zs_dim + zsa_hdim, zsa_hdim)   
13 = Linear(state_dim + action_dim + zs_dim + zsa_hdim * 2, zsa_hdim)   
...   
l_num_layers = Linear(state_dim + action_dim + zs_dim + zsa_hdim * (num_layers-1),zsa_hdim)   
final_layer = Linear(state_dim + action_dim + zs_dim + zsa_dim, target_dim)   
State-Action Encoder  $f^{sa}$  Forward Pass:   
input = concatenate([action,zs])   
x = swish(batchnorm(11(input)))   
input = concatenate([input,x])   
x = swish(batchnorm(12(input)))   
...   
x = swish(batchnorm(1_num_layers(input)))   
output = concatenate([input,x])   
Final Layer t Forward Pass:   
input = zsa   
output = final_layer(zsa)
C.2 SAC & TQC Hyperparameters
Our SAC baseline [Haarnoja et al., 2018] is based on our TD3 implementation, keeping hyperparameters constant when possible. Remaining details are based on the author implementation https://github.com/haarnoja/sac. Our TQC baseline [Kuznetsov et al., 2020] uses the author PyTorch implementation https://github.com/SamsungLabs/tqcPytorch (with evaluation code kept consistent for all methods). For hyperparameters, the action space is assumed to be in the range  $[-1,1]$ .
TQC is based on SAC and uses similar hyperparameters. One exception is the use of an  $\epsilon$  offset. As suggested in the appendix of Haarnoja et al. [2018], the log probability of the Tanh Normal distribution is calculated as follows:
$$
\log \pi (a | s) = \log N (u | s) - \log (1 - \tanh  (u) ^ {2} + \epsilon), \tag {33}
$$
where  $u$  is the pre-activation value of the action, sampled from  $N$ , where  $N$  is distributed according to  $\mathcal{N}(\mu, \sigma^2)$ , from the outputs  $\mu$  and  $\log \sigma$  of the actor network. Kuznetsov et al. [2020] use an alternate calculation of the log of the Tanh Normal distribution which eliminates the need for an  $\epsilon$  offset:
$$
\log \pi (a | s) = \log N (u | s) - (2 \log (2) + \log (\operatorname {s i g m o i d} (2 u)) + \log (\operatorname {s i g m o i d} (- 2 u))). \tag {34}
$$
Environment-specific hyperparameters. TQC varies the number of dropped quantiles depending on the environment. While it is possible to select this quantity based on heuristics [Kuznetsov et al., 2021], we use the author suggested values for each environment.
Hyperparameters of SAC are listed in Table 5. Network architecture details of SAC are described in Pseudocode 5.
Hyperparameters of TQC are listed in Table 6. Network architecture details of TQC are described in Pseudocode 6.
Table 5: SAC Hyperparameters.
Hyperparameter	Value
SAC
[Haarnoja et al., 2018]	Target Entropy	-action_dim
Policy log standard deviation clamp	[-20, 2]
Numerical stability offset ε	1e-6
Exploration	Initial random exploration time steps	25k
Common	Discount factor γ	0.99
Replay buffer capacity	1M
Mini-batch size	256
Target update rate τ	0.005
Optimizer	(Shared) Optimizer	Adam [Kingma and Ba, 2014]
(Shared) Learning rate	3e-4
Pseudocode 5. SAC Network Details
Value  $Q$  Network:
$\triangleright$  SAC uses two value networks each with the same network and forward pass.
l1 = Linear(state_dim + action_dim, 256)
$12 =$  Linear(256，256)
$l3 =$  Linear(256,1)
Value  $Q$  Forward Pass:
input = concatenate(state, action])
x = ReLU(11(input))
$x = \operatorname {ReLU}(12(x))$
value  $= 13(x)$
Policy  $\pi$  Network:
11 = Linear(state_dim, 256)
12 = Linear(256, 256)
13 = Linear(256, action_dim * 2)
Policy  $\pi$  Forward Pass:
input  $=$  state
$\mathbf{x} = \mathrm{ReLU}(11(\mathrm{input}))$
$x = \operatorname {ReLU}(12(x))$
mean,  $\log_{-}\mathrm{std} = 13(x)$
x = Normal(mean, exp(log_stdclip(-20, 2))).sample()
action  $=$  tanh(x)
Table 6: TQC Hyperparameters.
Hyperparameter	Value
TQC[Kuznetsov et al., 2020]	Number of networks	5
Quantiles per network	25
Quantiles dropped per network
HalfCheetah	0
Hopper	5
Walker2d	2
Ant	2
Humanoid	2
SAC[Haarnoja et al., 2018]	Target Entropy	-action_dim
Policy log standard deviation clamp	[-20,2]
Exploration	Initial random exploration time steps	256
Common	Discount factor γ	0.99
Replay buffer capacity	1M
Mini-batch size	256
Target update rate τ	0.005
Optimizer	(Shared) Optimizer	Adam [Kingma and Ba, 2014]
(Shared) Learning rate	3e-4
Pseudocode 6. TQC Network Details
Value  $Q$  Network:
$\triangleright$  TQC uses five value networks each with the same network and forward pass.
Value  $Q$  Forward Pass:
l1 = Linear(state_dim + action_dim, 512)  
l2 = Linear(512, 512)  
l3 = Linear(512, 512)  
l4 = Linear(512, 25)
input = concatenate(state, action)]  
x = ReLU(11(input))  
x = ReLU(12(x))  
x = ReLU(13(x))  
value = 14(x)
Policy  $\pi$  Network:
11 = Linear(state_dim, 256)  
12 = Linear(256, 256)  
13 = Linear(256, action_dim * 2)  
Policy π Forward Pass:  
input = state  
x = ReLU(11(input))  
x = ReLU(12(x))  
mean, log_std = 13(x)  
x = Normal(mean, exp(log_std.clip(-20, 2))).sample()  
action = tanh(x)
C.3 Offline RL Baselines
We use four offline RL baseline methods. The results of each method are obtained by re-running the author-provided code with the following commands:
CQL [Kumar et al., 2020b]. https://github.com/aviralkumar2907/CQL commit d67dbe9
python examples/cq1_mujoco_new.py --env=ENV \
--policy_1r=3e-5 --langrange thresh=-1 --min_q_weight=10
Setting the Lagrange threshold below 0 means the Lagrange version of the code is not used (as dictated by the settings defined in the CQL paper). Default hyperparameters in the GitHub performed substantially worse (total normalized score:  $221.3 \pm 34.4$ ). Some modifications were made to the evaluation code to match the evaluation protocol used by TD7.
TD3+BC [Fujimoto and Gu, 2021]. https://github.com/sfujim/TD3_BC commit 8791ad7
python main.py --env=ENV
IQL [Kostrikov et al., 2021]. https://github.com/ikostrikov/implicit_q_learningcommit09d7002
python train_offline.py --env_name  $\equiv$  ENV --config=configs/mujoco_config.py
$\mathcal{X}$  -QL [Garg et al., 2023]. https://github.com/Div99/XQL commit dff09af
python train_offline.py --env_name  $\equiv$  ENV --config=configs/mujoco_config.py   
--max_clip  $= 7$  --double  $\equiv$  True --temp  $= 2$  --batch_size  $= 256$
This setting corresponds to the  $\mathcal{X}$ -QL-C version of the method where hyperparameters are fixed across datasets.
D Design Study
In this section we analyze the space of design choices for SALE. All results are based on modifying TD7 without policy checkpoints. A summary of all results is contained in Table 7. Each choice is described in detail in the subsections that follow.
Table 7: 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
Learning Target (Section D.1)
s'	16438 ± 219	3507 ± 313	5882 ± 262	7415 ± 753	6337 ± 387
z's'	17317 ± 65	3034 ± 647	6227 ± 218	7939 ± 334	6984 ± 578
z's and r	17089 ± 219	3264 ± 670	6255 ± 220	8300 ± 675	6841 ± 252
z's'a'	12445 ± 417	3009 ± 528	4757 ± 352	6288 ± 570	5657 ± 110
Network Input (Section D.2)
Q, remove zsa	16948 ± 250	2914 ± 756	5746 ± 254	6911 ± 552	6907 ± 443
Q, remove zs	16856 ± 241	2783 ± 710	5987 ± 407	7954 ± 248	7196 ± 363
Q, remove s,a	17042 ± 208	2268 ± 723	5486 ± 466	8546 ± 655	6449 ± 478
Q, zsa only	15410 ± 691	2781 ± 585	5205 ± 409	7916 ± 976	5551 ± 1144
Q, s,a only	12776 ± 774	2929 ± 631	4938 ± 422	5695 ± 875	6209 ± 384
π, s only	16997 ± 265	3206 ± 419	5744 ± 430	7710 ± 593	6413 ± 1521
π, z's only	17283 ± 388	2429 ± 705	6319 ± 331	8378 ± 732	6507 ± 496
No fixed embeddings	17116 ± 149	2262 ± 590	5835 ± 481	8027 ± 551	7334 ± 153
Normalization (Section D.3)
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
End-to-end (Section D.4)
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
D.1 Learning Target
The encoder  $(f_{t + 1},g_{t + 1})$  in TD7 uses the following update:
$$
\begin{array}{l} \mathcal {L} \left(f _ {t + 1}, g _ {t + 1}\right) := \left(g _ {t + 1} \left(f _ {t + 1} (s), a\right) - \left| f _ {t + 1} \left(s ^ {\prime}\right) \right| _ {\times}\right) ^ {2} (35) \\ = \left(z _ {t + 1} ^ {s a} - \left| z _ {t + 1} ^ {s ^ {\prime}} \right| _ {\times}\right) ^ {2}. (36) \\ \end{array}
$$
In this section, we vary the learning target (originally  $z_{t + 1}^{s^{\prime}}$ ).
(Figure 7) Target 1:  $s'$ . Sets the learning target to the next state  $s'$ , making the encoder loss:
$$
\mathcal {L} \left(f _ {t + 1}, g _ {t + 1}\right) := \left(z _ {t + 1} ^ {s a} - s ^ {\prime}\right) ^ {2}. \tag {37}
$$
This target is inspired by OFENet [Ota et al., 2020].
(Figure 7) Target 2:  $z_{\text{target}}^{s'}$ . Sets the learning target to the embedding of the next state  $z_{\text{target}}^{s'}$  taken from a slow-moving target network. This makes the encoder loss:
$$
\mathcal {L} \left(f _ {t + 1}, g _ {t + 1}\right) := \left(z _ {t + 1} ^ {s a} - \left| z _ {t + 1} ^ {s a} \right| _ {\times}\right) ^ {2}. \tag {38}
$$
$z_{\mathrm{target}}^{s^{\prime}}$  is taken from a target network  $f_{\mathrm{target}}$  where  $f_{t + 1}$  has parameters  $\theta$  and  $f_{\mathrm{target}}$  has parameters  $\theta^{\prime}$ , and  $\theta^{\prime}$  is updated at each time step by:
$$
\theta^ {\prime} \leftarrow (0. 9 9) \theta^ {\prime} + (0. 0 1) \theta . \tag {39}
$$
This target is inspired by SPR [Schwarzer et al., 2020].
(Figure 7) Target 3:  $z^s$  and  $r$ . Sets the learning target to the embedding of the next state  $z_{t+1}^s$  (which is the same as in the original loss) but adds a second loss towards the target  $r$ . This is achieved by having the output of  $g$  be two dimensional, with the first dimension being  $z_{t+1}^{sa}$  and the second  $r_{t+1}^{\mathrm{pred}}$ , making the encoder loss:
$$
\begin{array}{l} \mathcal {L} \left(f _ {t + 1}, g _ {t + 1}\right) := \left(g _ {t + 1} \left(f _ {t + 1} (s), a\right) [ 0 ] - \left| f _ {t + 1} \left(s ^ {\prime}\right) \right| _ {\times}\right) ^ {2} + \left(g _ {t + 1} \left(f _ {t + 1} (s), a\right) [ 1 ] - r\right) ^ {2} (40) \\ = \left(z _ {t + 1} ^ {s a} - \left| z _ {t + 1} ^ {s ^ {\prime}} \right| _ {\times}\right) ^ {2} + \left(r _ {t + 1} ^ {\text {p r e d}} - r\right) ^ {2}. (41) \\ \end{array}
$$
This target is inspired by DeepMDP [Gelada et al., 2019].
(Figure 8) Target 4:  $z^{s'a'}$ . Sets the learning target to the next state-action embedding  $z_{t+1}^{s'a'}$ , where the action  $a'$  is sampled according to the target policy with added noise and clipping (as done in TD3 [Fujimoto et al., 2018]). This makes the encoder loss:
$$
\mathcal {L} \left(f _ {t + 1}, g _ {t + 1}\right) := \left(z _ {t + 1} ^ {s a} - \left| z _ {t + 1} ^ {s ^ {\prime} a ^ {\prime}} \right| _ {\times}\right) ^ {2}, \tag {42}
$$
$$
a ^ {\prime} := \pi_ {t} \left(z _ {t - 1} ^ {s ^ {\prime}}, s ^ {\prime}\right) + \epsilon , \tag {43}
$$
$$
\epsilon \sim \operatorname {c l i p} \left(\mathcal {N} \left(0, \sigma^ {2}\right), - c, c\right). \tag {44}
$$
Table 8: 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
s'	16438 ± 219	3507 ± 313	5882 ± 262	7415 ± 753	6337 ± 387
z's'	17317 ± 65	3034 ± 647	6227 ± 218	7939 ± 334	6984 ± 578
z's and r	17089 ± 219	3264 ± 670	6255 ± 220	8300 ± 675	6841 ± 252
z's/a'	12445 ± 417	3009 ± 528	4757 ± 352	6288 ± 570	5657 ± 110
Figure 7: Learning targets from baseline methods. Learning curves on the MuJoCo benchmark, using learning targets inspired by baseline methods, the next state  $s'$  from OFENet [Ota et al., 2020], the next state embedding  $z_{\text{target}}^{s'}$  from a slow-moving target network from SPR [Schwarzer et al., 2020], and including the reward term  $(z^{s'}$  and  $r)$  as in DeepMDP [Gelada et al., 2019]. Note that this study only accounts for changes to the learning target and does not encompass all aspects of the baseline algorithms. Results are averaged over 10 seeds. The shaded area captures a  $95\%$  confidence interval around the average performance.
Figure 8: Alternate learning target. Learning curves on the MuJoCo benchmark, using a learning target of the next state-action embedding  $z^{s'a'}$ . Results are averaged over 10 seeds. The shaded area captures a  $95\%$  confidence interval around the average performance.
D.2 Network Input
SALE modifies the value function and policy to take the embeddings as additional input:
$$
Q _ {t + 1} (s, a) \rightarrow Q _ {t + 1} \left(z _ {t} ^ {s a}, z _ {t} ^ {s}, s, a\right), \quad \pi_ {t + 1} (s) \rightarrow \pi_ {t + 1} \left(z ^ {s}, s\right). \tag {45}
$$
In this section, we vary the input of the value function and policy. In all cases, changes are made independently (i.e. we fix the policy to the above (using embeddings) and change the policy, or fix the value function to the above (using embeddings) and change the policy).
(Figure 9) Input 1:  $Q$ , remove  $z^{sa}$ . We remove  $z^{sa}$  from the value function input. This makes the input:  $Q_{t + 1}(z_t^s, s, a)$ .
(Figure 9) Input 2:  $Q$ , remove  $z^s$ . We remove  $z^s$  from the value function input. This makes the input:  $Q_{t + 1}(z_t^{sa}, s, a)$ .
(Figure 9) Input 3:  $Q$ , remove  $s, a$ . We remove  $s, a$  from the value function input. This makes the input:  $Q_{t + 1}(z_t^{sa}, z^s)$ .
(Figure 10) Input 4:  $Q$ ,  $z^{sa}$  only. We only use  $z^{sa}$  in the value function input. This makes the input:  $Q_{t+1}(z_t^{sa})$ .
(Figure 10) Input 5:  $Q, s, a$  only. We only use  $s, a$  in the value function input. This makes the input:  $Q_{t+1}(s, a)$ . As for all cases, the policy still takes in the state embedding  $z^s$ .
(Figure 11) Input 6:  $\pi$ ,  $s$  only. We remove  $z^s$  from the policy input, using only the state  $s$  as input. This makes the input:  $\pi_{t+1}(s)$ . Note that the value function still uses the embeddings as input.
(Figure 11) Input 7:  $\pi, z^s$  only. We remove  $s$  from the policy input, using only the embedding  $z^s$  as input. This makes the input:  $\pi_{t+1}(z_t^s)$ .
(Figure 12) Input 8: No fixed embeddings. We remove the fixed embeddings from the value function and the policy. This means that the networks use embeddings from the current encoder, rather than the fixed encoder from the previous iteration. This makes the input:  $Q_{t + 1}(z_{t + 1}^{sa}, z_{t + 1}^s, s, a)$  and  $\pi_{t + 1}(z_{t + 1}^s, s)$ .
Table 9: 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
Q, remove zsa	16948 ± 250	2914 ± 756	5746 ± 254	6911 ± 552	6907 ± 443
Q, remove zs	16856 ± 241	2783 ± 710	5987 ± 407	7954 ± 248	7196 ± 363
Q, remove s,a	17042 ± 208	2268 ± 723	5486 ± 466	8546 ± 655	6449 ± 478
Q,zsaonly	15410 ± 691	2781 ± 585	5205 ± 409	7916 ± 976	5551 ± 1144
Q,s,a only	12776 ± 774	2929 ± 631	4938 ± 422	5695 ± 875	6209 ± 384
π,s only	16997 ± 265	3206 ± 419	5744 ± 430	7710 ± 593	6413 ± 1521
π,zsonly	17283 ± 388	2429 ± 705	6319 ± 331	8378 ± 732	6507 ± 496
No fixed embeddings	17116 ± 149	2262 ± 590	5835 ± 481	8027 ± 551	7334 ± 153