PaperlayerView exact revision on arXiv

Source-linked AI summary

Representation Matters: Offline Pretraining for Sequential Decision Making

Mengjiao Yang, Ofir Nachum

arXiv:2102.05815v1cs.LGcs.AI

TL;DR

Offline RL research has largely emphasized policy optimization from static data; this paper asks whether offline experience can instead pretrain state representations for downstream decision-making. It evaluates unsupervised representation objectives across imitation learning, offline RL, and online RL, finding broad gains and task-dependent objective strengths.

  • Problem

    Offline RL has largely focused on policy optimization, leaving the question of whether offline experience can pretrain state representations that improve downstream policy learning.

    Source: §1 · p. 1

  • Method

    The paper evaluates unsupervised objectives on offline datasets by pretraining state representations for imitation learning, offline RL, and online RL tasks.

    Source: §1 · pp. 1–2

  • Results

    Simple representation-learning objectives improve downstream imitation learning, offline RL, and online RL, but the ideal objective appears to depend on the task.

    Source: §1 · p. 3, §6 · p. 11

  • Takeaways & Limitations

    Contrastive self-prediction objectives are promising, and the ablations indicate that representation-learning choices should be considered in relation to the downstream setting.

    Source: §1 · p. 3, §5.1 · p. 11, §6 · p. 11

    The study focuses on state representation learning and leaves combining it with action representation learning for future work.

    Source: §1 · p. 2

Abstract

from arXiv · show

The recent success of supervised learning methods on ever larger offline datasets has spurred interest in the reinforcement learning (RL) field to investigate whether the same paradigms can be translated to RL algorithms. This research area, known as offline RL, has largely focused on offline policy optimization, aiming to find a return-maximizing policy exclusively from offline data. In this paper, we consider a slightly different approach to incorporating offline data into sequential decision-making. We aim to answer the question, what unsupervised objectives applied to offline datasets are able to learn state representations which elevate performance on downstream tasks, whether those downstream tasks be online RL, imitation learning from expert demonstrations, or even offline policy optimization based on the same offline dataset? Through a variety of experiments utilizing standard offline RL datasets, we find that the use of pretraining with unsupervised learning objectives can dramatically improve the performance of policy learning algorithms that otherwise yield mediocre performance on their own. Extensive ablations further provide insights into what components of these unsupervised objectives -- e.g., reward prediction, continuous or discrete representations, pretraining or finetuning -- are most important and in which settings.

1 Introduction

The paper asks whether offline experience can pretrain state representations that improve downstream policy learning, rather than focusing only on offline policy optimization. It tests diverse objectives across imitation learning, offline RL, and online RL, finding that contrastive self-prediction objectives are especially promising.

1 Introduction

  • The paper asks whether representations pretrained on offline experience can improve downstream policy learning from state observations, unlike prior work focused on behavior or action representations.The offline-RL literature described here largely emphasizes policy optimization from static datasets.Source: §1 · pp. 1–2
  • Experiments use D4RL Gym-MuJoCo datasets and evaluate low-data imitation, offline RL on the pretraining dataset, and online RL with environment access.The three downstream settings distinguish small expert demonstrations, the same offline data, and online interaction, respectively.Source: §1 · p. 2
  • The study focuses on state representations and leaves combining state and action representations to future work.The authors describe this boundary as a choice to avoid an overly broad paper.Source: §1 · p. 2
  • Contrastive self-prediction objectives are the promising group: they predict components of a sub-trajectory from other components, using a contrastive loss for state prediction.An example predicts the next state from a sequence of states and actions.Source: §1 · p. 3
  • Simple representation learning objectives on offline data can improve and accelerate downstream imitation learning, offline RL, and online RL compared with learning from raw observations.The paper highlights imitation from far-from-expert pretraining data and offline RL where pretraining and downstream data are the same.Source: §1 · p. 3

2 Background and Related Work

Related work spans state abstraction, model-based latent representations, and contrastive objectives, while most prior approaches described here learn representations online. This paper instead studies offline pretraining across distinct downstream settings, including non-image tasks.

Abstraction and Bisimulation

  • Traditional abstraction and bisimulation methods seek compact state or action spaces by grouping states and actions with similar rewards and dynamics.More recent neural approaches learn stochastic representations and reward or transition functions in embedding space.Source: §2 · p. 3

Representations in Model-Based Learning

  • Model-based approaches use latent representations to simplify model learning and policy rollouts, while inverse-model approaches predict actions from preceding and subsequent state representations.Latent-space rollouts are sometimes called imagination.Source: §2 · p. 3

Contrastive Objectives

  • Contrastive representation methods encourage similar states to have similar representations, with similarity often based on temporal proximity or observation augmentations.Prior work connects contrastive loss with mutual information maximization and energy-based models.Source: §2 · pp. 3–4

State Representation Learning in Offline RL

  • Most related work described focuses on online representation learning, whereas this study separates representation learning on static offline data from downstream task learning.The paper notes that static data makes method comparisons easier by avoiding exploration and nonstationary-dataset confounds.Source: §2 · p. 4
  • The paper distinguishes its study through broad comparisons across settings and benefits on non-image tasks without image-specific prior knowledge.Prior offline-pretraining studies mentioned here typically initialize image encoders, often using expert demonstrations.Source: §2 · p. 4

3 Task Setups

The experimental protocol evaluates offline-pretrained representations in imitation learning, offline RL, and online RL using Gym-MuJoCo data. It specifies how datasets and downstream tasks are constructed and how results are measured.

3 Task Setups

  • The protocol covers low-data imitation, offline policy optimization, and online RL in partially observable environments, with 16 data-task combinations per setting and five seeds.The setup totals 240 training runs for each representation-learning objective.Source: Table 1 · p. 4, §3 · p. 4

3.1 Datasets

  • Gym-MuJoCo datasets cover halfcheetah, hopper, walker2d, and ant, with expert, medium-expert, medium, and medium-replay data reflecting collection-policy quality.Datasets contain trajectories, and datasets generated by multiple policies need not be Markovian even when the underlying environment is Markovian.Source: §3.1 · pp. 4–5

3.2 Imitation Learning in Low-Data Regime

  • Imitation learning tests behavioral cloning on the first 10k or 25k transitions of expert data, using medium or medium-replay data from the same domain for representation learning.The protocol tests whether representations from larger mediocre-behavior datasets help cloning from a small expert dataset.Source: §3.2 · p. 5

3.3 Offline RL with Behavior Regularization

  • Offline RL uses the same selected domain-and-dataset combination for representation pretraining and downstream BRAC policy learning.BRAC uses a KL-divergence behavior penalty; the study compares learned representations with raw states, including performance and hyperparameter dependence.Source: §3.3 · p. 5

3.4 Online RL in Partially Observable Environments

  • Online RL applies SAC downstream with offline-pretrained representations in partially observable domains created by randomly zero-masking a state dimension.The same masking is applied to the offline dataset used for pretraining.Source: §3.4 · pp. 5–6

3.5 Evaluation

  • Each run uses five seeds, 200k pretraining steps, and 1M downstream training steps, with periodic evaluation and normalized returns averaged over the final ten evaluations.Pretraining batches contain 256 sampled sub-trajectories; policy evaluation uses ten episodes, and final scores aggregate across seeds and setups.Source: §3.5 · p. 6

4 Experiments: Breadth Study

Across downstream tasks, representation pretraining improves average performance, with contrastive self-prediction objectives among the strongest. Longer sub-trajectories generally perform better, while forward-model objectives perform poorly in these settings.

4.1 Representation Learning Objectives

  • ACL masks entries in state-action-reward sub-trajectories and uses a transformer to contrast masked-state representations with outputs at the corresponding sequence positions.The implementation uses k + 1 = 8.Source: §4.1 · p. 7

4.2 Results

  • Pretraining improves downstream imitation learning, offline RL, and online RL by 1.5x, 2.5x, and 15%, respectively.Figure 2 aggregates average rewards over five seeds; methods that failed to converge are excluded.Source: Figure 2 · p. 8, §4.2 · p. 8
  • ACL, (Momentum) TCL, and VPN are among the strongest objectives; they use contrastive self-prediction, predicting sub-trajectory components from others and applying contrastive loss when predicting states.ACL performs well on imitation learning, while VPN and (momentum) TCL perform well on offline and online RL.Source: §4.2 · p. 8, Figure 2 · p. 8
  • Sub-trajectories of length k + 1 = 8 generally outperform k + 1 = 2, possibly because datasets can be non-Markovian, including when distinct policies collect data.Source: §4.2 · p. 8
  • Raw-state prediction and future-representation objectives (DeepMDP and Bisimulation) performed poorly in these settings; extensive tuning and related attempts yielded no better results, making these baselines unattractive versus simpler, better-performing alternatives.Source: §4.2 · p. 8

5 Experiments: Depth Study

The depth study varies components of ACL-based contrastive self-prediction and evaluates their effects across imitation learning, offline RL, and online RL. The results show that useful choices differ by downstream task, with auxiliary losses particularly mixed across settings.

5 Experiments: Depth Study

  • The study uses transformer-based ACL as an ablation skeleton, changing one factor at a time from its default configuration; the pretrained representation is reused downstream unless context embedding is enabled.Input action and input reward default to true; other listed factors default to false.Source: Figure 3 · p. 9, Table 2 · p. 9, §5 · p. 10

5.1 Results

  • Action and reward inclusion, prediction, or reconstruction benefits offline and online RL but impairs imitation learning; no single objective dominates all settings.Reconstruction is reported as a more consistent benefit in RL settings than including or predicting actions and rewards.Source: Figure 4 · p. 10, §5.1 · p. 11
  • Context embedding consistently hurts performance; bidirectional transformers hurt imitation but help online RL, while finetuning modestly degrades imitation and offline RL and can help online RL by domain.Source: §5.1 · p. 11
  • Using the representation objective as an auxiliary loss can dramatically improve online RL but dramatically degrade imitation learning and offline RL.The auxiliary-loss approach is also used in online RL literature.Source: §5.1 · p. 11

6 Conclusion

Simple representation-learning objectives can dramatically improve downstream imitation learning, offline RL, and online RL; the best objective may depend on the task, and paradigms popular in NLP or online RL may perform poorly offline.

6 Conclusion

  • Simple representation-learning objectives can dramatically improve downstream imitation learning, offline RL, and online RL, but the ideal objective may depend on the task.The ablations also suggest that paradigms popular in NLP or online RL may not translate to good performance in offline settings.Source: §6 · p. 11
  • The study leaves multi-task, transfer, and exploration tasks, real-world domains such as robot learning, and network-architecture ablations for future work.These are identified as aspects beyond the scope of the current study.Source: §6 · p. 11

A.1 Representation Network

  • The representation network is a two-hidden-layer fully connected network with 256 units per layer; dimensions 256 and 512 generally work best among those tested.The tested output dimensions were 16, 64, 256, and 512.Source: §A.1 · p. 15

A.2 Transformer Network

  • ACL's BERT-style transformer uses four 128-unit attention heads; masking drops items with probability 0.3, switches them with probability 0.15, or keeps them with probability 0.15.The masking rates were chosen arbitrarily and not tuned.Source: §A.2 · p. 15

A.3 Action Prediction and Reconstruction

  • Action prediction or reconstruction uses a tanh-squashed Gaussian output distribution and an additive adaptive entropy regularizer on the prediction loss.The implementation follows Haarnoja et al. (2019).Source: §A.3 · p. 15

A.4 Other Networks

  • Most other functions f and g are two-hidden-layer fully connected networks with 256 units per layer and Swish activations; Momentum TCL adds a residual connection to f's output.The Momentum TCL exception retains the same network structure apart from that residual connection.Source: §A.4 · p. 15

A.5 Training

Pretraining generally uses Adam, with a higher learning rate for TCL variants. Several objectives diverged on offline and online RL, while bisimulation sometimes failed to converge.

A.5 Training

  • Pretraining uses Adam at learning rate 0.0001, except TCL variants use 0.0003; Momentum TCL uses a moving-average rate of 0.05.Source: §A.5 · p. 15

A.6 Convergence Failures

  • Forward-raw model, VPN, and DeepMDP representations consistently diverged and output NaNs on offline and online RL; bisimulation failed in some runs.The divergent objectives were excluded from Figure 2, while bisimulation means were computed from successful runs only.Source: §A.6 · p. 15

B.1 Ablation Over Representation Size

Representation dimensions 256 and 512 work best in the reported ablation.

B.1 Ablation Over Representation Size

  • Dimensions 256 and 512 work best across domains and datasets when action and reward reconstruction are enabled.Source: Figure 5 · p. 16

B.2 Additional Contrastive Learning Results

The best ACL variant generally performs best among the contrastive objectives, while TCL remains competitive in offline RL.

B.2 Additional Contrastive Learning Results

  • The best ACL variant generally performs best against other contrastive objectives, although TCL is competitive in offline RL.The compared ACL settings vary by downstream task: input reward is off for imitation learning, action reconstruction is on for offline RL, and auxiliary loss or finetuning is on for online RL.Source: Figure 6 · p. 16

B.3 Ablation Results for Individual Domains and Datasets

Across the individual-domain ablations, representation learning has benefits in each imitation-learning combination, offline-RL benefits are clearer with expert trajectories, and online-RL effects vary with finetuning.

B.3 Ablation Results for Individual Domains and Datasets

  • Some ACL variant outperforms learning without representation in every imitation-learning domain-dataset combination.Inputting or reconstructing actions and rewards is especially harmful in several domains.Source: Figure 7 · p. 17
  • Offline-RL representation benefits are more evident with expert trajectories, while action and reward reconstruction matters more in ant and halfcheetah.The expert-present datasets include expert and medium-expert; expert-absent examples include medium and medium-replay.Source: Figure 8 · p. 18
  • Auxiliary loss generally improves online-RL performance, but finetuning improves halfcheetah, hopper, and walker2d while impairing ant.Source: Figure 9 · p. 19

C Additional Anecdotal Conclusions

Additional ablations found no changed overall conclusions from combined factor changes, no significant effects from varying transformer block or head counts, and no benefits from normalization or stochastic regularization; action-only reconstruction and dimension-token transformers showed more qualified promise.

C Additional Anecdotal Conclusions

Related papers

Loading 2102.05815v1…