For SALE: State-Action Representation Learning for Deep Reinforcement Learning — part 2/4

Cleaned-text code points [15812, 26194). Bibliography omitted here; source.md includes it.

Clipped Values. Extrapolation error is the tendency for deep value functions to extrapolate to unrealistic values on state-actions pairs which are rarely seen in the dataset ( fujimoto2019off ) . Extrapolation error has a significant impact in offline RL, where the RL agent learns from a given dataset rather than collecting its own experience, as the lack of feedback on overestimated values can result in divergence.

Surprisingly, we observe a similar phenomenon in online RL, when increasing the number of dimensions in the state-action input to the value function, as illustrated in Figure 2 . Our hypothesis is that the state-action embedding z s ​ a z^{sa} expands the action input and makes the value function more likely to over-extrapolate on unknown actions. We show in Figure 2 that the dimension size of z s ​ a z^{sa} as well as the state-action input plays an important role in the stability of value estimates.

Fortunately, extrapolation error can be combated in a straightforward manner in online RL, where poor estimates are corrected by feedback from interacting with the environment. Consequently, we only need to stabilize the value estimate until the correction occurs. This can be achieved in SALE by tracking the range of values in the dataset D D (estimated over sampled mini-batches during training), and then bounding the target used in Equation 6 by the range: Q t + 1 ​ ( s , a ) ≈ r + γ ​ clip ​ ( Q t ​ ( s ′ , a ′ ) , min ( s , a ) ∈ D ⁡ Q t ​ ( s , a ) , max ( s , a ) ∈ D ⁡ Q t ​ ( s , a ) ) . \displaystyle Q_{t+1}(s,a)\approx r+\gamma\text{ clip}\left(Q_{t}(s^{\prime},a^{\prime}),\min_{(s,a)\in D}Q_{t}(s,a),\max_{(s,a)\in D}Q_{t}(s,a)\right). (9) Additional discussion of extrapolation error, experimental details, and ablation of the proposed value clipping in SALE can be found in & .

### 4.2 Evaluating Design Choices

The effectiveness of learning embeddings by modeling the dynamics of the environment is a natural consequence of the relationship between the value function and future states. However, there are many design considerations for which all alternatives are potentially valid and the approach adopted differs among related methods in the literature. In this section, we perform an extensive study over the design space to (1) show SALE uses the correct and highest performing set of choices, and (2) better understand which choices are the biggest contributors to performance when using SALE.

In Figure 3 we display the mean percent loss when modifying SALE in the TD7 algorithm (to be fully introduced in Section 5.2 ). The percent loss is determined from the average performance at 1M time steps, over 10 seeds and five benchmark environments (HalfCheetah, Hopper, Walker2d, Ant, Humanoid) ( OpenAIGym ) . A more detailed description of each variation and complete learning curves can be found in .

Learning target. TD7 trains the encoders by minimizing the MSE between the state-action embedding z s ​ a z^{sa} and a learning target of the next state embedding z s ′ z^{s^{\prime}} ( Equation 2 ). We test several alternate learning targets. OFENet uses the next state s ′ s^{\prime} as the target ( ota2020can ) while SPR ( schwarzer2020data ) uses the embedding z target s ′ z_{\text{target}}^{s^{\prime}} from a target network obtained with an exponential moving average with weight 0.01 0.01 . Drawing inspiration from Bisimulation metrics ( ferns2011bisimulation ) , DeepMDP ( gelada2019deepmdp ) use an objective that considers both the next state embedding z s ′ z^{s^{\prime}} and the reward r r . We test including a prediction loss on the reward by having the encoder g g output both z s ​ a z^{sa} and r pred r^{\text{pred}} where r pred r^{\text{pred}} is trained with the MSE to the reward r r . Finally, we test the next state-action embedding z s ′ ​ a ′ z^{s^{\prime}a^{\prime}} as the target, where the action a ′ a^{\prime} is sampled from the target policy.

Network input. In our approach, the learned embeddings ( z s ​ a , z s ) (z^{sa},z^{s}) are appended to the state and action input to the value function Q ⁡ ( z s ​ a , z s , s , a ) Q(z^{sa},z^{s},s,a) and policy π ⁡ ( z s , s ) \pi(z^{s},s) ( Equation 3 ). We attempt different combinations of input to both networks. We also evaluate replacing the fixed embeddings (Equations 6 & 7 ), with the non-static current embeddings ( z t + 1 s ​ a , z t + 1 s ) (z^{sa}_{t+1},z^{s}_{t+1}) .

Normalization. TD7 uses AvgL1Norm ( Equation 4 ) to normalize the scale of the state embedding z s z^{s} , as well as on the state-action input ( s , a ) (s,a) , following a linear layer ( Equation 5 ). We attempt removing AvgL1Norm on ( s , a ) (s,a) , removing it entirely, and adding it to the state-action embedding z s ​ a z^{sa} . We additionally test swapping AvgL1Norm for BatchNorm ( ioffe2015batch ) and LayerNorm ( ba2016layer ) . Finally, instead of directly applying normalization to the embeddings, we replace the MSE in the encoder loss ( Equation 2 ) by the cosine loss from schwarzer2020data .

End-to-end. Embeddings can be trained independently or end-to-end with the downstream task. We test our approach as an auxiliary loss to the value function. The encoders and the value function are trained end-to-end, thus allowing the value loss to affect the embeddings ( z s ​ a , z s ) (z^{sa},z^{s}) , where the encoder loss ( Equation 2 ) is multiplied by a constant to weigh its importance versus the value loss.

## 5 Stabilizing RL with Decoupled Representation Learning

In this section, we present the TD7 algorithm (TD3 + 4 additions). We begin by introducing the use of checkpoints in RL to improve the stability of RL agents. We then combine SALE with checkpoints and various previous algorithmic modifications to TD3 ( fujimoto2018addressing ) to create a single RL algorithm for both the online and offline setting.

### 5.1 Policy Checkpoints

Deep RL algorithms are notoriously unstable ( henderson2017deep ) . The unreliable nature of deep RL algorithms suggest a need for stabilizing techniques. While we can often directly address the source of instability, some amount of instability is inherent to the combination of function approximation and RL. In this section, we propose the use of checkpoints, to preserve evaluation performance, irrespective of the quality of the current learned policy.

A checkpoint is a snapshot of the parameters of a model, captured at a specific time during training. In supervised learning, checkpoints are often used to recall a previous set of high-performing parameters based on validation error, and maintain a consistent performance across evaluations ( vaswani2017attention ; kenton2019bert ) . Yet this technique is surprisingly absent from the deep RL toolkit for stabilizing policy performance.

In RL, using the checkpoint of a policy that obtained a high reward during training, instead of the current policy, could improve the stability of the performance at test time.

For off-policy deep RL algorithms, the standard training paradigm is to train after each time step (typically at a one-to-one ratio: one gradient step for one data point). However, this means that the policy changes throughout each episode, making it hard to evaluate the performance. Similar to many on-policy algorithms ( williams1992reinforce ; PPO ) , we propose to keep the policy fixed for several assessment episodes, then batch the training that would have occurred. • Standard off-policy RL: Collect a data point → \rightarrow train once.

• Proposed: Collect N N data points over several assessment episodes → \rightarrow train N N times.

In a similar manner to evolutionary approaches ( salimans2017evolution ) , we can use these assessment episodes to judge if the current policy outperforms the previous best policy and checkpoint accordingly. At evaluation time, the checkpoint policy is used, rather than the current policy.

We make two additional modifications to this basic strategy.

Minimum over mean. Setting aside practical considerations, the optimal approach would be to evaluate the average performance of each policy using as many trials as possible. However, to preserve learning speed and sample efficiency, it is only sensible to use a handful of trials. As such, to penalize unstable policies using a finite number of assessment episodes, we use the minimum performance, rather than the mean performance. This approach also means that extra assessment episodes do not need to be wasted on poorly performing policies, since training can resume early if the performance of any episode falls below the checkpoint performance.

Variable assessment length. In , we examine the caliber of policies trained with a varied number of assessment episodes and observe that a surprisingly high number of episodes (20 + ) can be used without compromising the performance of the final policy. However, the use of many assessment episodes negatively impacts the early performance of the agent. Freezing training for many episodes means that the environment is explored by a stale policy, reducing data diversity, and delaying feedback from policy updates. To counteract this effect, we restrict the number of assessment episodes used during the initial phase of training before increasing it.

Additional details of our approach to policy checkpoints can be found in .

### 5.2 TD7

TD7 is based on TD3 ( fujimoto2018addressing ) with LAP ( fujimoto2020equivalence ) , a behavior cloning term for offline RL ( fujimoto2021minimalist ) , SALE ( Section 4.1 ), and policy checkpoints ( Section 5.1 ).

LAP. Gathered experience is stored in a replay buffer ( expreplay1992 ) and sampled according to LAP ( fujimoto2020equivalence ) , a prioritized replay buffer D D ( PrioritizedExpReplay ) where a transition tuple i := ( s , a , r , s ′ ) i:=(s,a,r,s^{\prime}) is sampled with probability p ⁡ ( i ) = max ⁡ ( | δ ⁡ ( i ) | α , 1 ) ∑ j ∈ D max ⁡ ( | δ ⁡ ( j ) | α , 1 ) , where ​ δ ​ ( i ) := Q ⁡ ( s , a ) − y , \displaystyle p(i)=\frac{\max\left(|\delta(i)|^{\alpha},1\right)}{\sum_{j\in D}\max\left(|\delta(j)|^{\alpha},1\right)},\qquad\text{ where }\delta(i):=Q(s,a)-y, (10) where y y is the learning target. The amount of prioritization used is controlled by a hyperparameter α \alpha . Furthermore, the value function loss uses the Huber loss ( huber1964robust ) , rather than the MSE.