PaperlayerView exact revision on arXiv

Source-linked AI summary

For SALE: State-Action Representation Learning for Deep Reinforcement Learning

Scott Fujimoto, Wei-Di Chang, Edward J. Smith, Shixiang Shane Gu, Doina Precup, David Meger

arXiv:2306.02451v2cs.LGcs.AIstat.ML

TL;DR

Reinforcement learning is sample inefficient, and representation learning has largely been used for image-based rather than low-level-state tasks. This paper introduces SALE's joint state-action embeddings and combines SALE with checkpoints in TD7, which outperforms existing continuous-control methods in online and offline settings.

  • Problem

    Reinforcement learning is sample inefficient, while representation learning has mainly been applied to image-based tasks rather than environments with low-level states.

    Source: §1 · p. 1, §7 · p. 10

  • Method

    SALE learns joint state-action embeddings by modeling latent environment dynamics, and TD7 combines SALE and policy checkpoints with TD3 and other improvements.

    Source: §1 · pp. 1–2, §7 · p. 10

  • Results

    TD7 significantly outperforms state-of-the-art continuous-control methods in final performance and early learning, and matches expensive offline algorithms.

    Source: §7 · p. 10

  • Takeaways & Limitations

    SALE and policy checkpoints are general-purpose techniques that can be included with most RL methods.

    Source: §7 · p. 10

    Because most representation-learning methods do not cover low-level states, the benchmark compares TD7 with only one other method focused on them.

    Source: J Limitations · p. 52

Abstract

from arXiv · show

In the field of reinforcement learning (RL), representation learning is a proven tool for complex image-based tasks, but is often overlooked for environments with low-level states, such as physical control problems. This paper introduces SALE, a novel approach for learning embeddings that model the nuanced interaction between state and action, enabling effective representation learning from low-level states. We extensively study the design space of these embeddings and highlight important design considerations. We integrate SALE and an adaptation of checkpoints for RL into TD3 to form the TD7 algorithm, which significantly outperforms existing continuous control algorithms. On OpenAI gym benchmark tasks, TD7 has an average performance gain of 276.7% and 50.7% over TD3 at 300k and 5M time steps, respectively, and works in both the online and offline settings.

1 Introduction

The introduction motivates learning state-action representations for low-level states and presents SALE, a latent-dynamics approach. It also previews design evaluation, clipping to mitigate extrapolation error, checkpoints, and the TD7 combination with TD3.

Motivation

  • RL is described as sample inefficient, while representation learning has mainly been applied to vision-based environments rather than low-level states.The paper suggests task difficulty can reflect the underlying dynamical system rather than observation-space size.Source: §1 · p. 1

SALE

  • SALE learns joint state-action embeddings by modeling environment dynamics in latent space, extending prior work.The introduction argues that meaningful features can come from state-action interactions even when the observation space is already compact.Source: §1 · p. 1
  • Expanding the action-dependent input reveals extrapolation error, which the authors mitigate with a simple clipping technique [Fujimoto et al., 2019].Source: §1 · p. 2

Design study

  • The paper evaluates SALE's design space to identify which choices contribute most to final performance.The introduction presents this as an empirical study because the best design choices are unclear.Source: §1 · p. 2

Checkpoints

  • At test time, checkpoints use the highest-performing policy observed during training after policies are fixed for multiple episodes.Source: §1 · p. 2

TD7

  • TD7 combines TD3 with SALE, checkpoints, prioritized experience replay [Fujimoto et al., 2020], and an offline-only behavior-cloning term.The introduction reports that TD7 significantly outperforms existing methods in online and offline RL without added complexity such as large ensembles, extra updates per time step, or per-environment hyperparameters.Source: §1 · p. 2

2 Related Work

Related work spans abstraction, compression, and feature learning, including dynamics-prediction approaches and action representations. The paper distinguishes SALE through joint state-action representations and its attention to design choices and decoupled-embedding difficulties.

Representation learning

  • Representation learning in RL has referred to abstraction, compression, and feature learning, and is commonly used for high-dimensional image observations.Source: §2 · p. 2
  • Predicting future states connects model-free and latent-space model-based RL, while another line learns representations over actions; SALE emphasizes joint state-action representations.Source: §2 · p. 2
  • Some state-action methods emphasize abstraction, whereas SALE extends OFENet with design considerations and remedies for decoupled-embedding difficulties.The paper reports that these improvements appear as significant benchmark performance gains.Source: §2 · p. 2

Stability in RL

  • Policy checkpoints are related to safe policy improvement and evolutionary RL; the checkpoint resembles the fittest individual, with mutation supplied by the underlying RL algorithm.Source: §2 · pp. 2–3

3 Background

The background frames reinforcement learning as an MDP in which a policy maps states to actions to maximize discounted reward. It also introduces the value function as modeling expected return from a state-action pair.

Background

  • An MDP is defined by state and action spaces, a reward function, dynamics, and a discount factor; its policy maps states to actions to maximize discounted reward.Source: §3 · p. 3
  • The value function Q^π(s, a) models expected return starting from state s and action a.Source: §3 · p. 3

4 State-Action Representation Learning

SALE learns state and state-action embeddings by predicting next-state structure, then supplies them alongside original inputs to the value and policy networks. The section examines stability measures and design choices for these representations.

4.1 State-Action Learned Embeddings

  • SALE encodes state as zs and state-action as zsa, training both encoders with an MSE prediction of the next-state embedding, independent of the next action or current policy.The embeddings are concatenated with the original state and action because they may omit information relevant to reward, policy, or task horizon.Source: §4.1 · p. 3
  • The encoders are trained online but decoupled from value-function and policy gradients; the embeddings improve network inputs rather than serve as a planning world model.The encoder objective uses stop-gradient on the next-state target.Source: §4.1 · p. 3
  • Increasing state-action input dimensions can destabilize online value estimates; clipping targets to observed value ranges stabilizes them without reducing input dimensions.The reported instability is associated with the state-action embedding and the linear state-action input, while the state embedding remains unchanged.Source: §4.1 · pp. 4–5, Figure 2 · p. 5

4.2 Evaluating Design Choices

  • Next-state-based learning targets perform similarly, with the next-state embedding improving performance; the next state-action target performs much worse, while reward prediction has little impact.The authors associate the weaker next state-action target with a non-stationary policy signal that can harm learning.Source: §4.2 · pp. 5–6
  • Learned features benefit both value and policy networks, but matter more for the value function; all components of the value-function input are needed for highest performance.Using zsa alone performs poorly, whereas combining it with the original state and action significantly improves performance.Source: §4.2 · p. 6
  • AvgL1Norm keeps embedding and downstream input scales controlled; in the design study, it outperforms alternate normalization approaches.The state-action embedding is not normalized because it is trained to match the normalized next-state embedding.Source: §4.1 · p. 4, §4.2 · p. 6
  • End-to-end training of embeddings with the value function performs significantly worse than decoupled representation learning.The tested end-to-end variant lets value-loss gradients affect the embeddings alongside the encoder loss.Source: §4.2 · p. 6

5 Stabilizing RL with Decoupled Representation Learning

The section introduces TD7, combining SALE and policy checkpoints with TD3 and other modifications. Checkpoints preserve evaluation performance by using a high-reward policy snapshot instead of relying on the current policy.

5 Stabilizing RL with Decoupled Representation Learning

  • TD7 combines SALE and policy checkpoints with TD3 and other algorithmic modifications for online and offline RL.The section also introduces LAP and, for offline RL, a behavior-cloning term.Source: §5 · p. 7, §5.2 · pp. 7–8

5.1 Policy Checkpoints

  • Policy checkpoints preserve evaluation performance by selecting a high-reward policy snapshot instead of relying on the current policy.The approach keeps the policy fixed for several assessment episodes, then batches the training that would otherwise have occurred.Source: §5.1 · p. 7
  • Checkpoint selection uses minimum episode performance rather than the mean to penalize unstable policies with few assessment trials.Training can resume early if any assessment episode falls below the checkpoint performance.Source: §5.1 · p. 7
  • Using many assessment episodes can hurt early performance by freezing training, leaving exploration to a stale policy and delaying update feedback.The number of assessment episodes is restricted early in training before being increased.Source: §5.1 · p. 7
  • Offline TD7 adds a behavior-cloning loss, but does not use checkpoints because it does not interact with the environment.The same loss function is used online and offline, with λ = 0 online.Source: §5.2 · p. 8

6 Results

On MuJoCo, TD7 significantly outperforms TQC and TD3+OFE at both early and final training, and often exceeds TD3’s 5M-step performance by 300k steps. On D4RL, TD7 outperforms all reported baselines; the section also describes ablation and runtime evaluations.

Online

  • The online MuJoCo benchmark compares TD7 with TD3, SAC, TQC, and TD3+OFE using a consistent evaluation protocol.SAC and TD3+OFE were re-implemented; TD3 and TQC used author-provided code.Source: §6 · p. 8
  • TD7 significantly outperforms TQC and TD3+OFE at 300k and 5M time steps, and often surpasses TD3’s 5M-step performance by 300k.The MuJoCo learning curves average 10 seeds with 95% confidence intervals; Table 1 reports average performance at 300k, 1M, and 5M time steps over 10 trials.Source: §6 · p. 8, Table 1 · p. 9, Figure 4 · p. 9

Ablation study and Run time

  • The study reports an ablation of SALE, checkpoints, and LAP, and benchmarks runtime across online baselines using identical resources and framework.The supplied passages describe these evaluations but do not report their outcomes.Source: §6 · p. 9

Offline

  • TD7 outperforms all baselines on the D4RL benchmark after 1M time steps.The comparison uses consistent-protocol reruns of author-provided baseline code with a single set of hyperparameters.Source: §6 · p. 9, Table 2 · p. 10

7 Conclusion

The conclusion presents SALE as a way to apply representation learning to low-level states by modeling state-action interactions in latent dynamics. It highlights clipping and checkpoints as stability measures and reports that TD7 significantly outperforms state-of-the-art continuous-control methods in early and final performance.

7 Conclusion

  • SALE learns state-action embeddings through a latent dynamics model, extending representation learning beyond image-based tasks.The conclusion reports an empirical investigation of SALE design choices.Source: §7 · p. 10
  • Clipping targets with seen values corrects instability associated with extrapolation error from increased state-action embedding input dimensions.The conclusion also identifies policy checkpoints as a source of added stability.Source: §7 · p. 10
  • TD7 combines SALE and policy checkpoints with TD3 and other recent improvements, and significantly outperforms state-of-the-art continuous-control methods in early and final performance.The conclusion also says TD7 can match expensive offline algorithms.Source: §7 · p. 10

A.1 Algorithm

  • TD7 uses paired value functions and target value functions, policy and target-policy networks, encoders, fixed encoders, and checkpoint networks.Its policy update uses a deterministic policy-gradient loss with a behavior-cloning term for offline RL; λ is 0 online.Source: §A.1 · p. 17
  • The encoder has a state network and a state-action network that each output an embedding, and is updated with a dedicated loss.The value functions consume SALE embeddings alongside state-action inputs.Source: §A.1 · p. 17

A.2 Hyperparameters

  • TD7 assumes actions in [−1, 1] and periodically updates target networks every 250 time steps rather than using an exponential moving average.The periodic update is necessary because TD7 uses fixed encoders.Source: §A.2 · p. 19

B Experimental Details

The appendix specifies benchmark protocols, baseline implementations, and experimental design details needed to interpret and reproduce the evaluations.

B Experimental Details

  • MuJoCo v4 experiments leave state, action, and reward spaces unchanged; methods gather 25k initial random interactions (256 for TQC), are evaluated every 5k interactions over 10 episodes, and repeated across 10 seeds.Source: B Experimental Details · p. 21

C.1 TD3 & TD3+OFE Hyperparameters

  • TD3 and TD3+OFE use stated network architectures and shared hyperparameters, while OFE layer counts vary by environment and Humanoid omits hard-to-predict force dimensions.OFE uses 292 of Humanoid’s 376 state dimensions; other tasks use the full state.Source: §C.1 · pp. 22–23

C.2 SAC & TQC Hyperparameters

  • SAC and TQC baselines use author implementations, with TQC’s dropped-quantile count set to author-recommended environment-specific values.TQC’s log-probability calculation avoids an epsilon offset through an alternate Tanh Normal calculation.Source: §C.2 · pp. 25, 27

C.3 Offline RL Baselines

  • Offline comparisons rerun four author-provided baselines; X-QL-C uses fixed hyperparameters across datasets.CQL’s default GitHub settings performed substantially worse, and evaluation code was adjusted to match TD7’s protocol.Source: §C.3 · p. 28

D Design Study

  • The design study modifies TD7 without policy checkpoints and reports average MuJoCo benchmark performance at 1M steps over 10 seeds.The study’s summary is Table 7.Source: D Design Study · p. 29, Table 7 · p. 29

D.1 Learning Target

The section tests design choices for SALE’s learning targets, network inputs, normalization, and end-to-end training, then discusses extrapolation error associated with expanded action-dependent inputs.

D.1 Learning Target

  • Learning-target variants predict the next state, its embedding, state embedding plus reward, or the next state-action embedding using a noisy, clipped target-policy action.The embedding target may come from a slow-moving target network; the reward-augmented variant is inspired by DeepMDP, while other variants draw on OFENet and SPR.Source: §D.1 · p. 30

D.2 Network Input

  • Value-function inputs are varied by removing zsa, zs, or raw state-action inputs, using zsa alone, or retaining only the raw state-action pair.Changes to the policy and value-function inputs are tested independently.Source: §D.2 · p. 32
  • Policy-input variants remove either the state embedding or raw state; another variant replaces fixed embeddings with outputs from the current encoder.The fixed-embedding variant applies to both the value function and policy.Source: §D.2 · p. 32, Figure 12 · p. 33

D.3 Normalization

  • SALE applies AvgL1Norm to state embeddings and to linear-layer outputs used as value-function and policy inputs.The study also tests no normalization, normalization on zsa, and BatchNorm or LayerNorm in place of AvgL1Norm.Source: §D.3 · p. 34
  • A further normalization variant removes normalization and replaces the encoder’s MSE loss with cosine similarity, following SPR.Normalization is otherwise not used in that variant.Source: §D.3 · p. 34, Figure 16 · p. 36

D.4 End-to-end

  • The end-to-end variant trains encoders with the value-function loss rather than keeping the embeddings fixed, and varies the loss-weight hyperparameter β.Its performance is reported in Table 11 and Figure 17.Source: §D.4 · p. 37, Table 11 · p. 37, Figure 17 · p. 37

E.1 Extrapolation Error

  • The authors observe extrapolation error with SALE and hypothesize that expanding action-dependent input dimensions makes unseen actions more distinct from seen actions.For Ant, the action dimension is 8 while zsa has dimension 256; this proposed explanation is framed as a hypothesis.Source: §E.1 · p. 38

E.2 Empirical Analysis

The section varies value-function inputs and replay-buffer size to examine extrapolation error, motivates checkpoints with evidence of training instability, and describes a checkpoint strategy based on assessment performance.

E.2 Empirical Analysis

  • The empirical analysis varies value-function inputs, including clipping, zsa, the dimension of ϕ, and whether encoders are trained, while also varying replay-buffer capacity.The comparisons include no clipping, no zsa, smaller ϕ, both changes together, frozen embeddings, and TD7 without checkpoints.Source: §E.2 · p. 38
  • Ant plots show each seed’s performance and value estimate for replay buffers of 1M or 100k across the tested variants.The captions describe plotted conditions but do not establish which variant performs best.Source: Figure 18 · p. 39, Figure 19 · p. 39, Figure 20 · p. 39

F.1 Motivation

  • Averaged, smoothed TD3 curves can conceal instability; in the last 1M steps, roughly 30% of evaluations are followed by at least one evaluation 50% worse within 10 evaluations.The single-seed comparison contrasts standard smoothing and averaging with an unsmoothed single-seed curve.Source: §F.1 · p. 41, Figure 24 · p. 41, Figure 25 · p. 41

F.2 Method

  • Checkpointing assesses the current policy during training and updates the stored policy when it outperforms the checkpoint, using minimum performance to favor consistency.Assessment can terminate early when current minimum performance falls below the checkpoint’s minimum, avoiding further assessment of a suboptimal policy.Source: §F.2 · p. 41
  • Early learning uses one assessment episode, later increasing the count; the checkpoint performance is reduced by a factor of 0.9 when that increase occurs.The adjustment addresses the observed negative effect of long assessment phases on early learning and possible overfitting to a single assessment episode.Source: §F.2 · p. 42

F.3 Empirical Analysis

The empirical analysis varies policy-checkpoint assessment choices and compares TD7 with versions that remove its main components. It reports checkpoint results at 5M steps, immediate-checkpoint early-learning results at 300k, and individual-seed curves with and without checkpoints.

F.3 Empirical Analysis

  • The checkpoint analysis varies assessment length, minimum versus mean performance, early termination, and whether assessment length starts at one or remains fixed.The default uses the minimum over 20 assessment episodes, terminates assessment early when current performance falls below the checkpoint, and increases the maximum from one to 20 episodes after 750k steps.Source: §F.3 · p. 43
  • At 5M steps, Table 12 presents the main checkpoint results; Table 13 reports immediate-checkpoint performance at 300k, and Figure 30 shows five individual seeds with and without checkpoints.The tables average results over 10 seeds; Figure 30 shows learning curves with a 95% confidence interval.Source: §F.3 · p. 43, Table 12 · p. 43, Table 13 · p. 44

G.1 Main Components

  • The main-component ablation removes SALE, checkpoints, or LAP, and also tests a SALE encoder trained end-to-end with the value function.The no-LAP variant uses uniform replay sampling and mean-squared error instead of Huber loss.Source: §G.1 · p. 46, Figure 31 · p. 46

G.2 Subcomponents

The subcomponent study varies SALE encoders, checkpoints, LAP, clipping, normalization, fixed encoders, and implementation choices. The supplied passages also identify offline learning curves and runtime-adjusted curves as evaluations, without specifying their outcomes.

G.2 Subcomponents

  • The subcomponent ablations test removing normalization or fixed embeddings, and compare SALE encoders trained end-to-end with the value function against TD3 with an encoder.In the no-fixed-encoder variant, the networks use embeddings from the current encoder rather than the previous iteration's fixed encoder.Source: §G.2 · p. 47, Figure 36 · p. 49
  • Checkpoint variants compare TD7 using the current policy at test time with TD7 using a checkpoint policy, and add checkpoint policies to TD3.The corresponding learning-curve figure averages results over 10 seeds and shows a 95% confidence interval.Source: §G.2 · p. 47, Figure 33 · p. 48
  • The LAP and clipping ablations add LAP to TD3, remove TD7's value clipping, or add value clipping to TD3.The clipping is proposed to mitigate extrapolation error.Source: §G.2 · p. 47, Figure 34 · p. 49, Figure 35 · p. 49
  • Implementation ablations remove TD7's implementation differences or apply those differences to TD3, including use of both value functions for policy updates and ELU value-network activations.TD3 with these differences also updates its target network every 250 steps rather than using an exponential moving average.Source: §G.2 · p. 47, Figure 37 · p. 49

H Offline RL Learning Curves

  • Figure 38 presents learning curves on the offline D4RL benchmark, averaged over 10 seeds with a 95% confidence interval.The supplied passage identifies the benchmark and curve scope but does not state comparative outcomes.Source: Figure 38 · p. 50

I Run time

  • Runtime-adjusted MuJoCo curves plot performance against runtime over 5M steps, while the runtime benchmark excludes evaluation episodes and uses shared hardware and PyTorch.The benchmark uses one Nvidia Titan X GPU and an Intel Core i7-7700k CPU.Source: Figure 39 · p. 51, I Run time · p. 51

J Limitations

The paper identifies limits in benchmark breadth, representation-learning baselines, computational cost, and theoretical analysis.

J Limitations

  • The empirical study emphasizes depth over breadth and leaves other continuous-control benchmarks, including image-based observations, for future examination.The stated boundary concerns the scope of experimentation.Source: J Limitations · p. 52
  • TD7 is compared with one other method focused on representation learning, because most such methods do not cover low-level states.The authors say that testing additional methods could provide new insight into representation learning for low-level states.Source: J Limitations · p. 52
  • Although TD7 has lower computational cost than competing methods, its runtime is more than double TD3's.Source: J Limitations · p. 52
  • The paper reports extensive empirical analysis but does not address theoretical questions about why dynamics-based representation learning works or what makes a good representation.Source: J Limitations · p. 52

Related papers

Loading 2306.02451v2…