PaperlayerView exact revision on arXiv

Source-linked AI summary

Investigating the role of model-based learning in exploration and transfer

Jacob Walker, Eszter Vértes, Yazhe Li, Gabriel Dulac-Arnold, Ankesh Anand, Théophane Weber, Jessica B. Hamrick

arXiv:2302.04009v1cs.LG

TL;DR

Prior work had shown model-based learning can aid generalization in some cases but offered little insight into when or why. This paper studies intrinsic-exploration pretraining and downstream fine-tuning, finding model-based methods outperform model-free baselines.

  • Problem

    Prior work showed model-based learning can aid generalization in some cases but offered little insight into when it should help or why.

    Source: §1 · p. 1

  • Method

    The study compares model-based and model-free agents pretrained through intrinsic exploration, then fine-tuned across transfer settings with component-retention ablations.

    Source: §1 · pp. 1–2

  • Results

    Model-based optimization dominates model-free variants during both unsupervised exploration and fine-tuning, with transferred dynamics and prior policy contributing most to transfer.

    Source: §5 · p. 8

  • Takeaways & Limitations

    Environment shift between unsupervised pretraining and downstream phases is detrimental to knowledge transfer.

    Source: §5 · p. 8

    The study fixes its intrinsic exploration mechanism and does not compare alternative intrinsic rewards.

    Source: §5 · p. 8

Abstract

from arXiv · show

State of the art reinforcement learning has enabled training agents on tasks of ever increasing complexity. However, the current paradigm tends to favor training agents from scratch on every new task or on collections of tasks with a view towards generalizing to novel task configurations. The former suffers from poor data efficiency while the latter is difficult when test tasks are out-of-distribution. Agents that can effectively transfer their knowledge about the world pose a potential solution to these issues. In this paper, we investigate transfer learning in the context of model-based agents. Specifically, we aim to understand when exactly environment models have an advantage and why. We find that a model-based approach outperforms controlled model-free baselines for transfer learning. Through ablations, we show that both the policy and dynamics model learnt through exploration matter for successful transfer. We demonstrate our results across three domains which vary in their requirements for transfer: in-distribution procedural (Crafter), in-distribution identical (RoboDesk), and out-of-distribution (Meta-World). Our results show that intrinsic exploration combined with environment models present a viable direction towards agents that are self-supervised and able to generalize to novel reward functions.

1. Introduction

The paper asks when and why model-based learning helps transfer, studying agents pretrained through intrinsic exploration and fine-tuned on downstream tasks with different components retained or reset.

1. Introduction

  • The study isolates transferable world knowledge from task knowledge by using intrinsic-exploration pretraining and resetting or retaining components during fine-tuning.Source: §1 · p. 1
  • The model-based agent uses a self-supervised MuZero variant, while both agent types use RND exploration across Crafter, RoboDesk, and Meta-World transfer settings.Source: §1 · pp. 1–2
  • Model-based exploration combined with model-based fine-tuning yields better transfer performance than model-free baselines.The paper uses component-transfer ablations to investigate how the model and policy contribute.Source: §1 · p. 2
  • Model-based methods perform better exploration than model-free counterparts in reward-free environments.Source: §1 · p. 2
  • Knowledge transfers most effectively when both pretraining and fine-tuning are model-based.Source: §1 · p. 2
  • World-model dynamics appear to improve transfer, with the model-based advantage stronger when the dynamics model is trained on the same environment.Source: §1 · p. 2

2. Background

This background reviews ways to transfer knowledge in reinforcement learning and approaches to exploring without task rewards.

2.1. Transfer learning in RL

  • Transfer learning can reuse data, skills, behaviors, value functions, or dynamics, while unsupervised RL pretrains through intrinsic exploration before task fine-tuning.Source: §2.1 · p. 2
  • Unlike some prior model-transfer approaches, this study does not assume previously collected experience can be relabeled with the target reward.Source: §2.1 · pp. 2–3

2.2. Intrinsic exploration

  • Intrinsic exploration seeks general-purpose environmental exploration without an explicit task reward, including count-based, entropy-based, curiosity-driven, and empowerment approaches.Source: §2.2 · p. 3
  • The study uses RND, a curiosity-based method that computes prediction error against a fixed randomly initialized neural-network target.Source: §2.2 · p. 3

3. Methods

The methods pre-train agents through intrinsic exploration, then fine-tune transferred components on downstream task rewards. The design compares model-based and model-free learning across phases and tests transfer in environments with different degrees of shift.

3. Methods

  • Intrinsic-reward pre-training followed by task-reward-only fine-tuning is compared across four model-based/model-free phase combinations and component-reset ablations.The ablations examine dynamics heads, the dynamics model, prior heads, and the prior policy head separately from prior reward and value heads.Source: §3 · p. 3

3.1. Agents

  • MuZero with a self-predictive loss supplies the model-based agent, while a matched Q-learning architecture enables transfer of shared encoder and prior-head weights.MuZero learns reward, action, and value predictions and uses its model for planning; the Q-learning baseline shares the observation encoder and prior-head architecture.Source: §3.1 · pp. 3–4

3.2. Exploration

  • RND prediction error provides intrinsic exploration rewards, and using the dynamics model's output to construct its prediction produced inferior exploration.For MuZero, the RND signal is predicted rather than measured; the RND feature projection is integrated into the agent's loss.Source: §3.2 · p. 4

3.3. Environments

  • Transfer is tested in procedurally generated in-distribution Crafter, reward-shift-only RoboDesk, and out-of-distribution Meta-World with novel objects and configurations.Fine-tuning is limited to 1M environment steps, but unsupervised exploration has no data-consumption limit; the study focuses on relative transfer rather than benchmarking exploration techniques. RoboDesk uses pixel observations and sparse rewards, while Meta-World uses corner3 pixels and sparse success rates.Source: §3.3 · pp. 4–5

4. Results

Across Crafter and RoboDesk, model-based exploration and fine-tuning generally outperform model-free approaches. Transfer is less effective when downstream dynamics are out of distribution, as in Meta-World.

Q1

  • At 1M environment steps, MB→MB achieves a score of 16.4 ± 1.5 and reward of 12.7 ± 0.4 on Crafter, outperforming the other agent variations.It also improves on DreamerV3, a model-based agent without pre-training.Source: Figure 3 · p. 6, §4 · p. 6
  • During reward-free Crafter pre-training, model-based optimization produces more effective exploration than model-free optimization, measured by average and individual achievement success rates.The comparison is between agents maximizing the same RND intrinsic reward.Source: §4 · pp. 6–7
  • On RoboDesk, MB→MB consistently outperforms MF→MB across nine tasks; model-based warm-starting improves sample efficiency over scratch training but not final performance.Pre-training success rates did not correlate with positive transfer across tasks.Source: Figure 4 · p. 6, §4 · p. 7

Q2

  • Crafter ablations indicate that transferring the dynamics model and heads, especially the policy prior, contributes most to positive transfer.Removing dynamics heads causes a small drop; removing the model causes a further drop, while removing prior heads also reduces performance, mainly through the policy prior.Source: §4 · p. 7, Figure 5 · p. 7

Q3

  • On Meta-World ML-10, MB→MB outperforms other pre-trained baselines but has only a small advantage over random initialization early in training, up to ∼20 million frames.The authors speculate that out-of-distribution dynamics limit transfer because only robotic-arm dynamics are shared between training and test environments.Source: §4 · pp. 7–8, Figure 6 · p. 8

5. Discussion

The discussion frames transfer as moving world knowledge learned without rewards into tasks with new reward functions. The results favor model-based optimization, while environment shift and fixed exploration choices bound the study's scope.

5. Discussion

  • The proposed framework studies transfer from unsupervised, reward-free exploration into task-specific fine-tuning, carrying representations, policies, and dynamics models between phases.Experiments span procedural in-distribution, identical in-distribution, and out-of-distribution environments.Source: §5 · p. 8
  • Model-based optimization outperforms model-free variants in exploration and fine-tuning; transferring dynamics and heads, especially the prior policy, contributes most to transfer.Environment shift between the unsupervised and downstream phases is detrimental to knowledge transfer.Source: §5 · p. 8
  • The study fixes the intrinsic exploration mechanism, so its comparisons address model-based versus model-free approaches rather than alternative intrinsic rewards.The authors also expect broader qualitative conclusions and suggest sample-efficient base agents or empowerment as possible alternatives.Source: §5 · p. 8

A.1. Network Architecture

The model-free and model-based agents use the same architecture, with input resolution and downsampling varying by environment. The encoder and dynamics model are ResNets with 10 blocks.

A.1. Network Architecture

  • Crafter and Meta-World use 96 × 96 pixel inputs, while RoboDesk uses 64 × 64 inputs.The architecture is identical for the model-free and model-based agents and matches MuZero ReAnalyse.Source: §A.1 · p. 12
  • The convolutional stack downsamples inputs to a 6 × 6 tensor for Meta-World or an 8 × 8 tensor for RoboDesk and Crafter before the encoder.Source: §A.1 · p. 12
  • The encoder and dynamics model each use a 10-block ResNet, with two convolutional 3 × 3 layers and 256 planes per block.Source: §A.1 · p. 12

A.2. Model-based Agent

The model-based agent adds an SPR objective that trains its dynamics predictions to match projections of future observations.

A.2. Model-based Agent

  • SPR matches each dynamics-model prediction vector to a target projection of the corresponding future observation using cosine distance.Target projections use the target network and replace the image-history stack with 15 copies of the current image.Source: §A.2 · p. 12
  • The projection and prediction heads use matching two-convolution architectures, with the projector receiving dynamics output and the predictor receiving projector output.Source: §A.2 · p. 12
  • The dynamics model unrolls n steps to produce prediction vectors matched against target vectors computed at the corresponding future steps.Source: §A.2 · p. 12
  • Table 1 lists hyperparameters for model-based pre-training and fine-tuning.Source: Table 1 · p. 12

A.3. Model-free Agent

The model-free baseline adapts the model-based agent’s architecture and training targets while removing model-based components.

A.3. Model-free Agent

  • The controlled Q-learning agent retains the model-based agent’s network architecture but modifies it to be model-free.Source: §A.3 · pp. 12–13
  • The Q-learning baseline uses n-step action-value targets, with its target network updated every 100 training steps.Source: §A.3 · p. 13
  • Its total loss separates reward and value losses, and the value target assumes the predicted immediate reward equals the true reward.The resulting target is used to compute the value loss.Source: §A.3 · p. 13
  • For continuous-action RoboDesk and Meta-World, the agent maximizes over sampled actions rather than the entire action space.Source: §A.3 · p. 13
  • The model-free agent unrolls a single step for computing actions and values.Source: §A.3 · p. 13
  • Table 2 lists hyperparameters for model-free pre-training and fine-tuning.Source: Table 2 · p. 13

A.4. Random Network Distillation

Random network distillation supplies intrinsic reward from prediction error while sharing the agent’s encoder.

A.4. Random Network Distillation

  • RND compares a learned prediction vector with a randomly initialized target vector, while the shared observation encoder receives both RND and agent-loss training signals.The prediction error is their L2 distance.Source: §A.4 · p. 13
  • Intrinsic reward is the prediction error normalized by its exponentially averaged mean and standard deviation, with no reward clipping.The moving average uses 0.99 decay and bias correction.Source: §A.4 · p. 14

B. Additional Results

On Crafter, model-based pre-training and fine-tuning have especially large success-rate margins on advanced tasks involving stone.

B.1. Crafter

  • MB→MB has particularly large success-rate margins on Crafter’s more advanced fine-tuned subtasks, especially those requiring or involving stone.Source: Figure 7 · p. 14, §B.1 · p. 14

B.2. RoboDesk

On RoboDesk, exploration success rates do not significantly correlate with task rewards, while return-based comparisons among transfer agents mirror the reported success-rate comparisons. On Meta-World, exploration likewise does not correlate with downstream rewards, and transfer settings perform similarly because transfer yields no tangible benefits.

B.2. RoboDesk

  • RoboDesk pre-training success rates show no significant correlation between exploration and task rewards.Figure 8 covers all 18 tasks.Source: §B.2 · p. 14, Figure 8 · p. 15
  • On RoboDesk, return-based comparisons among model-based, model-free, and from-scratch agents closely match the reported success-rate comparisons.Figure 9 reports fine-tuning returns on nine core tasks, using the median of three seeds and standard deviation.Source: §B.2 · p. 14, Figure 9 · p. 15

B.3. Meta-World

  • On Meta-World, exploration does not correlate with downstream task rewards, and transfer settings perform similarly because transfer provides no tangible benefits.Figure 10a shows exploration during pre-training; Figure 10b reports the fine-tuning ablation results.Source: §B.3 · p. 16, Figure 10 · p. 16
Loading 2302.04009v1…