Explore then Execute: Adapting without Rewards via Factorized Meta-Reinforcement Learning
Abstract
We seek to efﬁciently learn by leveraging shared structure between different tasks and environments.
For example, cooking is similar in different kitchens, even though the ingredients may change location.
In principle, meta-reinforcement learning approaches can exploit this shared structure, but in practice, they fail to adapt to new environments when adaptation requires targeted exploration (e.g., exploring the cabinets to ﬁnd ingredients in a new kitchen).
We show that existing approaches fail due to a chicken-and-egg problem: learning what to explore requires knowing what information is critical for solving the task, but learning to solve the task requires already gathering this information via exploration.
For example, exploring to ﬁnd the ingredients only helps a robot prepare a meal if it already knows how to cook, but the robot can only learn to cook if it already knows where the ingredients are.
To address this, we propose a new exploration objective (D REAM ), based on identifying key information in the environment, independent of this information will exactly be used solve the task.
By decoupling exploration from task execution,
D REAM explores and consequently adapts to new environments requiring no reward signal when the task is speciﬁed via an instruction.
Empirically, D REAM scales to more complex problems, such as sparse-reward 3D visual navigation, while existing approaches fail from insufﬁcient exploration.
1. Introduction
A general-purpose agent should be able to perform multiple tasks across multiple environments, both of which share considerable structure.
Our goal is to develop agents that can perform a variety of tasks in novel environments, based on previous experience and only a small amount of experience in the new environment.
For example, we may want a robot to cook a meal (a new task) in a new kitchen (the environment) after it has learned to cook other meals in other kitchens.
To adapt to a new kitchen, the robot must ﬁrst explore to discover the locations of the ingredients, and then use this information to cook.
Existing meta-RL methods have shown promise for adapting to new tasks and environments, but, as we identify in this work, are poorly equipped when adaptation requiring sophisticated exploration.
In the meta-RL setting, the agent is presented with a set of meta-training problems, each in an environment (e.g., a kitchen) with some task (e.g., make pizza); at meta-test time, the agent is given a new but similar environment and task.
It is allowed to gather experience in a few initial (training) episodes, and its goal is to then maximize returns on all subsequent (testing) episodes.
A common meta-RL approach is to train a recurrent neural network (RNN)
(Duan et al., 2016; Wang et al., 2016a; Stadie et al., 2018; Zintgraf et al., 2019; Humplik et al., 2019; Kamienny et al., 2020) to maximize test returns on each meta-training environment and task, such that it can quickly adapt to new meta-test settings.
With enough capacity, such approaches can express the optimal adaptation strategy, which involves both exploring in the initial few episodes (to ﬁnd relevant information) and then executing in the subsequent episodes (use the gained information to solve the task).
However, optimizing for exploration and execution end-to-end in this way creates a challenging chicken-and-egg optimization problem: Learning how to explore requires knowing what information is critical for actually solving the task, but learning to solve the task requires already gathering this information via exploration.
As a concrete example, exploring to locate the ingredients only helps a robot prepare a meal if it already knows how to cook, but the robot can only learn to cook if it already knows where the ingredients are.
The key insight of this work is to construct a decoupled exploration objective by ﬁrst automatically identifying taskrelevant information and then training an exploration policy to uncover this information.
Concretely, we assume access to a problem ID (during meta-training but not meta-testing) that identiﬁes the task and environment, which may contain information irrelevant to solving the task.
We ﬁrst train an execution policy to solve the tasks conditioned on an intermediate representation of this problem ID, and simultaneously minimize the mutual information between this
ID representation and the ID itself.
This makes the ID representation capture only task-relevant information, stripping extraneous information.
Then, we train an exploration policy to produce trajectories that contain precisely the information in the learned ID representation ensuring efﬁcient and targeted exploration (Section 4).
Beyond the above optimization challenges, a second issue prevents existing approaches from exploring to identify relevant information in the environment relates to the problem formulation rather than the approaches themselves.
In the standard meta-RL formulation, both the task and environment vary across problems and must be inferred via trialand-error.
While the environment must naturally be inferred through interaction, hiding all information about the task, e.g., the meal to cook (or the velocity to run at (Finn et al., 2017)) is unrealistic and can create unnecessary exploration challenges particularly in sparse reward settings.
For example, the robot might need to cook various meals until it guesses the correct meal to cook.
Instead, we propose a setting called instruction-based meta-RL (IMRL), where the agent is provided information about the task via instructions (e.g., ”cook a pizza” or a one-hot representation).
Providing instructions that contain information about the reward function also enables reward-free adaptation : the agent receives rewards during meta-training, but during meta-testing, it executes instructions in a new environment without ever receiving additional rewards, by exploring to identify relevant environment information and inferring the objective from the instructions (Section 3).
Overall, we present two contributions: (i) we identify the problem of coupling between exploration and execution and provide an approach that overcomes this via a decoupled objective, and (ii) a new meta-RL setting (IMRL), that can enable reward-free adaptation.
When applied to IMRL, our decoupled approach, called D REAM ( D ecoupling R ewardfree
E xplor
A tion and execution in M eta-RL) learns sophisticated exploration strategies and enables reward-free adaptation on a sparse-reward 3D visual navigation problem, achieving near-optimal performance.
In contrast, existing state-of-the-art meta-RL approaches (I MPORT , V ARI BAD,
RL 2 ) do not learn the optimal exploration strategy and achieve zero reward, even when they receive rewards at meta-test time (Section 6).
2. Preliminaries
Meta-Reinforcement Learning Setting.
The standard meta-RL setting considers a family of Markov decision processes (MDPs), (cid:104)S , A , R T , T T (cid:105) with states S , actions A , rewards
R T , and dynamics T T , parametrized by a problem
T .
Colloquially, we refer to the dynamics as the environment, the rewards as the task, and the entire MDP as the problem.
Different problems share structure (e.g., similar kitchens or cooking similar meals), and are drawn from some distribution T ∼ p ( T ) .
Following prior problem settings (Finn et al., 2017; Rakelly et al., 2019;
Rothfuss et al., 2018; Fakoor et al., 2019) and terminology from Duan et al.
(2016), we deﬁne a trial as M + N episodes in the same
MDP T , with M initial training (exploration) episodes of T steps (we set M = 1 for simplicity), typically used to infer dynamics and rewards, followed by N testing (execution) episodes.
To enable efﬁcient exploration, we allow the agent to take a special ”end-episode” action a term to terminate exploration early.
Meta-training and meta-testing both consist of sampling a problem T ∼ p ( T ) and running a trial.
The agent’s goal is to maximize the returns summed over the
N execution episodes during meta-testing.
Following prior work (Rakelly et al., 2019; Humplik et al., 2019; Kamienny et al., 2020), we make the problem ID of T available for meta-training, but not meta-testing, when the agent must infer this via exploration.
For convenience, we formally express the objective in terms of an exploration policy π exp used in the exploration episode and an execution policy π exe used in execution episodes, but these policies may be the same or share parameters.
The execution policy π exe may depend on all prior experiences in the trial, including the exploration trajectory τ exp =
( s 0 , a 0 , r 0 , . . .
, s T ) ∼ π exp from rolling out π exp in the exploration episode.
The goal of the agent to maximize:
[Equation]
where V π exe ( τ exp ; T ) is the expected return of π exe in problem T in the execution episodes.
Learning to Reinforcement Learn.
A common meta-RL approach is the Learning to Reinforcement Learn framework (Wang et al., 2016a; Duan et al., 2016; Zintgraf et al., 2019;
Kamienny et al., 2020; Humplik et al., 2019), which directly maximizes the objective J in (1) by learning the same recurrent policy π ( a t | s t , τ : t ) for both exploration and execution (i.e., π exe = π exp = π ).
This policy takes action a t given state s t and history of experiences spanning all episodes in a trial τ : t = ( s 0 , a 0 , r 0 , . . .
, s t − 1 , a t − 1 , r t − 1 ) .
Intuitively, the recurrent representation of τ : t can encode uncertainty over which problem the agent is currently in i.e., p ( T | τ : t ) (Kaelbling
[Table]
[Table] [Table]
The optimal policy explores to reduce this uncertainty in the initial exploration episode, ideally until the agent can receive high rewards via the same actions on all potential problems the agent might be in.
Then, when uncertainty is low during the execution episodes, the optimal policy can take actions leading to high returns.
By optimizing both exploration and execution end-to-end to maximize returns, this approach can learn the optimal policy, but optimization is challenging, which we show in Section 4.1.
We refer to this framework by its canonical instantiation, RL 2 (Duan et al., 2016).
3. Instruction-based Meta-RL
While the standard meta-RL setting (Section 2) captures adapting to new problems, it requires exploring to obtain both key information about the environment (the ingredients’ location in a kitchen), as well as the task itself (the meal to cook).
It is often more realistic and efﬁcient to directly convey the task to the agent.
We therefore propose a new meta-RL setting called instruction-based meta-RL (IMRL), where the agent receives instructions conveying the task.
Additionally, unlike the standard setting, where an agent repeatedly executes the same task (e.g., cook pizza) in a trial, we allow each episode to have a different task provided via a different instruction (e.g., cook pizza, then pasta, then soup), while the environment remains ﬁxed.
Speciﬁcally, we augment the state with an instruction i ∈
I (e.g., cook pizza), represented in this work as collections of onehots.
On each episode, we sample an instruction i from a distribution of instructions p T ( i ) and the agent receives the rewards R T ( s t , i ) at each timestep t in problem T .
Note that when the instruction is always the same, (e.g., |I| = 1 ),
IMRL can recover the standard setting, so any algorithms developed for IMRL (e.g., Section 4) also apply to the standard setting.
We summarize IMRL in Figure 1 for a special case below where we remove instructions and rewards from exploration episodes.
Special case: reward-free adaptation.
When the instructions i together with the transition observations, i.e., ( s, a, s (cid:48) ) -tuples, contain enough information about the reward function, executing instructions does not require observing the rewards directly during exploration.
For example, a robot can execute the instruction ”cook the soup recipe on the fridge” by reading (transition observations) and following the recipe on the fridge, without observing the corresponding rewards.
Consequently, we optionally make the exploration episode reward-free , providing neither instruction nor rewards.
Further, if the agent does not adapt using reward feedback during execution episodes, the entire process of adaptation at meta-test time can be made reward-free: Speciﬁcally, during meta-training, the exploration episode is reward-free, but the agent receives rewards during execution episodes.
During meta-testing, the agent receives no rewards at all.
This models real-world situations where providing rewards is expensive (e.g., a robot chef would ideally adapt to a new home kitchen without any human supervision).
IMRL allows us to conveniently model this special case.
4. Decoupling Exploration and Execution 4.1. The Problem with Coupling Exploration and Execution
We begin by showing how approaches like RL 2 , which optimize behavior on both exploration and execution episodes end-to-end by maximizing J in (1), can be sample inefﬁcient in IMRL, even when |I| = 1 .
For clarity, we refer to the policy followed during exploration episodes as π exp and the policy followed during execution episodes as π exe , although for RL 2 they are the same policy.
Figure 2a illustrates this end-to-end training approach.
Note that π exe relies on π exp for good exploration data in the form of τ exp in order to learn to solve the task.
Additionally, note that learning π exp relies on gradients passed through π exe .
If π exe cannot effectively solve the task, then these gradients will be uninformative.
This causes a bad local optimum: if our current (suboptimal) π exe obtains low rewards with a good informative trajectory ¯ τ exp , the low reward would cause π exp to learn to not generate ¯ τ exp .
This causes π exp to instead generate trajectories τ exp that lack information required to obtain high reward, further preventing the execution policy π exe from learning.
Typically, early in training, both π exp and π exe are stuck in this local optimum, where neither policy can become optimal without many samples.
We illustrate this effect in a simple example in Section 5.2.
4.2. D REAM : Decoupling Reward-free Exploration and Execution in Meta-learning
The key idea behind D REAM is to sidestep this local optimum by optimizing decoupled objectives for the exploration and execution policies.
Intuitively, the goal of the exploration policy is to identify the distinguishing characteristics of the environment and task relevant to executing instructions.
An initial approach is to train π exp so that it produces trajectories that are predictive of the problem ID or dynamics (Zhou et al., 2019b).
This yields a policy that can distinguish problems, but inefﬁciently explores instructionirrelevant attributes, such as the color of the walls.
To avoid this, we propose to derive a stochastic problem encoding
F ψ ( z | T ) , which extracts only the information necessary to execute instructions in the problem T .
D REAM obtains this encoder by training an execution policy π exe conditioned on the encoder’s outputs, with an information bottleneck on z .
Then, D REAM trains an exploration policy π exp to produce trajectories with high mutual information with z .
[Figure]
[Figure] [Figure]
[Figure] [Figure]
[Figure]
In this approach, the execution policy π exe no longer relies on effective exploration π exp to learn, and once F ψ ( z | T ) is learned, the the exploration policy π exp can also learn independent of π exe , decoupling the two optimization processes.
Later, during meta-testing, when the problem ID T is unavailable, the two policies easily combine, since the trajectories generated by π exp contain the same information as the stochastic encodings F ψ ( z | T ) that the execution policy π exe trained on (overview in Figure 2b).
Learning the problem ID encodings.
We begin with learning a good stochastic encoder F ψ ( z | T ) parametrized by ψ and execution policy π exe θ ( a | s, i, z ) parameterized by θ .
Unlike RL 2 , we choose not to make π exe θ recurrent for simplicity, relying on z to contain the necessary information.
We learn F ψ ( z | T ) jointly with the execution policy π exe θ ( a | s, i, z ) by optimizing the following objective:
[Equation]
Information bottleneck
[Equation]
where V π exe θ ( i, z ; T ) is the expected return of π exe θ on problem T given instruction i and encoding z .
Importantly, both terms are independent of the exploration policy π exp .
We minimize the mutual information I ( z ; T ) by applying a variational upper bound (Barber & Agakov, 2003) as follows.
[Equation]
where r is any prior and z is distributed as p ψ ( z )
=
[Equation]
T F
[Equation]
z
[Equation]
T
[Equation]
Learning an exploration policy given problem ID encodings.
Once we’ve obtained an encoder F ψ ( z | T ) to extract only the necessary information required to optimally execute instructions, we can optimize the exploration policy π exp to produce trajectories that encode this same information by maximizing their mutual information.
We slightly abuse notation to use π exp to denote the probability distribution over the trajectories τ exp .
Then, the mutual information I ( τ exp ; z ) can be efﬁciently optimized by applying a variational lower bound (Barber & Agakov, 2003) as follows.
[Equation]
where q ω is any distribution parameterized by ω .
We maximize the above expression over ω to learn q ω that approximates the true conditional distribution p ( z | τ exp ) , which makes this bound tight.
In addition we do not have access to the problem T at test time and hence cannot sample from
F ( z | T ) .
Therefore, q serves as a decoder that generates the encoding z from the exploration trajectory τ exp .
Note that only the third term depends on the exploration trajectory.
Hence, we maximize this by training the exploration policy on rewards set to be the information gain:
[Equation]
Intuitively, the reward for taking action a t and transitioning to state s t +1 is high if this transition encodes more information about the problem (and hence the encoding z ∼ F ψ ( z | T ) ) than was already present in the trajectory τ exp : t − 1 .
The reward also includes a small per timestep penalty c to encourage exploring efﬁciently.
This reward is attractive because (i) it is independent from the execution policy and hence avoids the local optima highlighted in Section 4.1, and (ii) it is a dense reward signal and helps with credit assignment.
Note that it is not Markovian, because it depends on τ exp , so we learn a recurrent π exp φ ( a t | s t , τ exp : t ) parametrized by φ that conditions on its history τ exp : t = ( s 0 , a 0 , r 0 , . . .
, s t − 1 , a t − 1 , r t − 1 ) .
For the reward-free exploration case, we simply omit the rewards r t from τ exp .
A Practical Implementation of D REAM
Altogether, D REAM learns four separate components (neural networks), which we detail below.
Encoder F ψ ( z | T ) : For simplicity, we parameterize the stochastic encoder by learning a deterministic encoding f ψ ( T ) and apply Gaussian noise, i.e., F ψ ( z | T ) =
N ( f ψ ( T ) , ρ 2 I ) .
We choose a convenient prior r ( z ) to be a unit Gaussian with same variance ρ 2 I which makes the information bottleneck (
Equation 3) take the form of
simple (cid:96) 2 -regularization (cid:107) f ψ ( T ) (cid:107) 22 . 2.
Decoder q ω ( z | τ exp )
: Similarly, we parameterize the decoder q ω ( z | τ exp ) as a Gaussian centered around a deterministic encoding g ω ( τ exp ) with variance ρ 2 I .
Then,
q ω maximizes E T ,z ∼ F ψ ( z |T ) (cid:104) (cid:107)
[Equation]
w.r.t., ω (
Equation 4),
and the
exploration
rewards take the form
[Equation] [Equation]
; T
[Equation]
−
[Equation]
(Equation 5).
3. Execution policy π exe θ and 4. Exploration policy π exp φ :
We represent both policies as Q-networks and apply double deep Q-learning (DDQN) (van Hasselt et al., 2016), treating ( s, i, z ) as the state for π exe θ .
While the above suggests training the encoder and execution policy to convergence, and then using these to train the exploration policy, we ﬁnd it less cumbersome to learn all components together in an EM-like fashion where in the exploration episode, we assume f ψ and π exe θ are ﬁxed.
We also train π exe θ conditioned on the exploration trajectory g ω ( τ exp ) , instead of exclusively training π exe on the outputs of the encoder z ∼ F ψ ( z | T ) .
We include all details and a summary (Algorithm 1) in Appendix A.
5. Analysis 5.1. Theoretical Consistency of the D REAM Objective
A key property of D REAM is that it is consistent : maximizing our decoupled objective also maximizes expected returns (Equation 1).
This contrasts prior works (Zhou et al., 2019b;
Rakelly et al., 2019; Gupta et al., 2018; Gurumurthy et al., 2019), which also decouple exploration from execution, but do not recover the optimal policy even with inﬁnite data.
Formally,
Proposition 1.
Let V ∗
( i ; T ) be the maximum expected returns achievable by any execution policy with access to T on instruction i and problem T , i.e., with complete information.
Let π exe (cid:63) , π exp (cid:63) , F (cid:63) and q (cid:63) ( z | τ exp ) be the optimizers of the D REAM objective.
Then for large enough T (the length of the exploration episode) and M (number of exploration episodes), and expressive enough function classes,
[Equation]
Since knowledge of T uniquely identiﬁes the MDP (cid:104)S , A , R T , T T (cid:105) , the maximum returns any agent can obtain corresponds to the returns of π exeopt , which has access to T .
Optimizing the decoupled objective in D REAM also obtains this maximum reward, which we prove in Appendix C.1.
5.2. Analysis of the Sample Complexity of Coupled and Decoupled Approaches
With enough capacity, approaches like RL 2 can also recover the optimal policy, but can be highly sample inefﬁcient due to the coupling problem in Section 4.1.
We highlight this in a simple tabular example to remove the effects of function approximation: Each episode is a one-step bandit problem with action space A .
Taking action a (cid:63) in the exploration episode leads to a trajectory τ exp (cid:63) that reveals the problem
T ; all other actions a reveal no information and lead to τ exp a .
The problem T identiﬁes a unique action that receives reward 1 during execution; all other actions get reward 0 .
Therefore, taking a (cid:63) during exploration is necessary and sufﬁcient to obtain optimal reward 1 .
We now study the number of samples required for RL 2 and D REAM to learn the optimal exploration policy with (cid:15) -greedy tabular Q-learning.
We precisely describe a more general setup in Appendix C. and prove that D REAM learns the optimal exploration policy in Ω( |A| H H log |A| ) times fewer samples than RL 2 in this simple setting with horizon H .
Figure 3a empirically validates this result and we provide intuition below.
In the tabular analog of RL 2 , the execution Q-values form targets for the exploration Q-values: ˆ Q exp ( a ) ← ˆ
V exe ( τ exp a ) := max a (cid:48)
ˆ Q exe ( τ exp a , a (cid:48) ) .
We drop the ﬁxed initial state from notation.
This creates the local optimum in Section 4.1.
Initially ˆ V exe ( τ exp (cid:63) ) is low, as the execution policy has not learned to achieve reward, even when given τ exp (cid:63) .
This causes ˆ Q exp ( a (cid:63) ) to be small and therefore arg max a ˆ Q exp ( a ) (cid:54) = a (cid:63) (Figure 3b), which then prevents ˆ
V exe ( τ exp (cid:63) ) from learning (Figure 3c) as τ exp (cid:63) is roughly sampled only once per |A| (cid:15) episodes.
This effect is mitigated only when ˆ Q exp ( a (cid:63) ) becomes higher than ˆ Q exp ( a ) for the other uninformative a s (the dot in Figure 3b-d).
Then, learning both the execution and exploration Q-values accelerates, but getting there takes many samples.
In D REAM , the exploration Q-values regress toward the decoder ˆ q : ˆ Q exp ( a ) ← log ˆ q ( T | τ exp ( a )) .
This decoder learns much faster than ˆ Q exe , since it does not depend on the execution actions.
Consequently, D REAM ’s exploration policy quickly becomes optimal (dot in Figure 3b), which allows quickly learning the execution Q-values and achieving high reward (Figures 3c and 3d).
6. Experiments
We compare D REAM with 4 state-of-the-art meta-RL algorithms: (i) RL 2 ; (ii) V ARI BAD (Zintgraf et al., 2019), which tries to optimally trades-off between exploration and exploitation with an auxiliary loss that predicts the dynamics and rewards from the recurrent state; (iii) I MPORT (Kamienny et al., 2020), which alternates between training the policy conditioned on the problem ID and its recurrent state and adds an auxiliary loss to keep the recurrent state close to the problem ID embedding; (iv) P EARL -UB, the analytically computed expected rewards achieved by the optimal problem-speciﬁc policy that explores via Thompson
Sampling (TS) with the true posterior problem distribution, which is an upperbound on the ﬁnal performance of
P EARL (Rakelly et al., 2019).
We report the average returns achieved by each method in trials with one reward-free exploration and one execution episode, averaged over 3 seeds with 1-std error bars (full details in Appendix B).
6.1. Didactic Experiments
We ﬁrst evaluate on grid world domains illustrated in Figures 4a and 4b.
The state is the agent’s ( x, y ) -position and a one-hot indicator of the object at the agent’s position (none, bus, map, pot, or fridge).
The actions are to move up, down, left, or right; ride bus (if the agent is at a bus, this teleports the agent to another bus of the same color); pickup ingredients at the fridge; and drop ingredients at the pot.
Episodes consist of 20 timesteps and the agent receives a reward of − 0 .
1 at each timestep until the instruction is executed (full details in Appendix B.1 and qualitative results in
Appendix B.2).
Targeted exploration.
We ﬁrst test if these methods can minimally explore to ﬁnd the necessary information with the family of problems in Figure 4a.
There are 4 possible instructions (green goal locations).
Reaching the correct goal yields +1 reward and ends the episode.
The four colored buses each lead to a different green goal location when ridden; in different problems T , their destinations are set to 1 of the 4!
different permutations.
In the distracting bus version, in different problems, the gray buses lead to different gray locations in the bottom row, which is unhelpful for solving the task.
In the map version, there is a map that reveals the destinations of the colored buses when touched.
Figures 5 and 6 (left) summarize the results.
D REAM learns to optimally explore and thus receives optimal reward in both versions: in distracting bus , it rides 3 of the 4 colored buses, and in map , it visits the map and ends the exploration episode.
During execution episodes, D REAM immediately reaches the goal by riding the correct colored bus.
In contrast, I MPORT and RL 2 rarely explore buses during exploration and then walk to the goal during execution, which achieves lower reward, indicative of the coupling problem (Section 4.1).
V ARI BAD learns slower, likely from learning the dynamics in its auxiliary loss, but eventually matches I MPORT and E-RL 2 in ˜3x as many samples.
The performance of P EARL -UB highlights how Thompson sampling even with the true posterior and optimal problem-speciﬁc policies does not explore optimally, and only explores one bus per episode (and never reads the map) in this case.
In distracting bus , we also compare to a version of DREAM without the information bottleneck (does not minimize the mutual information
I ( z ; T ) between problem IDs and the encoder F ψ ( z | T ) ).
As seen in Figure 5 (left), this ablation (D REAM (no bottleneck)) wastes its exploration on the gray unhelpful buses, since they are part of the problem, and consequently gets lower reward.
Generalization to new problems / instructions.
We test generalization to unseen problems / instructions in a cooking domain (Figure 4b).
The fridges on the right each contain 1 of 7 different (color-coded) ingredients, determined by the problem.
The fridges’ contents are unobserved until the agent uses the ”pickup” action at the fridge.
Instructions specify placing 2 correct ingredients in the pot in the right order.
We hold out ˜1% of the problems and ˜5% of the instructions during training.
Figure 6 shows the results on the training (left) and held-out (middle) problems and instructions.
Only D REAM achieves near-optimal returns on both.
During exploration, it investigates each fridge by using the ”pickup” action, and [Figure] [Figure] [Figure]
then directly retrieves the correct ingredients during execution.
RL 2 only sometimes explores the fridges during training, leading to lower returns.
It overﬁts to the training recipes and fails to cook the testing recipes.
Training on the problem ID actually hurts I MPORT compared to RL 2 .
I MPORT successfully cooks conditioned on the problem ID, but it fails on many recipes conditioned on its recurrent state and we observe the recurrent state and problem ID embeddings become far apart.
V ARI BAD learns slowly as before.
Thompson sampling (P EARL -UB) cannot learn optimal exploration as before.
6.2. Sparse-Reward 3D Visual Navigation
We conclude with a more challenging variant of the problem family introduced by Kamienny et al.
(2020) with visual observations, illustrated in (Figure 4c): There are two problems, one where the sign on the right says ”blue” and the other where it says ”red.”
The instructions specify whether the agent should go to the key, block, or ball, and the agent receives +1 reward for going to the object with the correct color (speciﬁed by the sign) and shape, and -1 reward otherwise.
The agent begins the episode on the far side of the barrier and must walk around the barrier to visually “read” the sign.
The agent’s observations are images and its actions are to rotate left or right, move forward, or end the episode.
D REAM is the only method that learns to read the sign and achieve reward (Figure 6 (right)).
RL 2 , I MPORT , and
V ARI BAD do not read the sign and consequently stay away from all the objects, in fear of receiving negative reward.
This occurs even when we allow these methods to receive an instruction and rewards during the exploration episode (dashed lines), while D REAM successfully adapts rewardfree .
Since the optimal problem-speciﬁc policy does not ever read the sign for any instruction or problem, optimal Thompson sampling exploration (P EARL -UB) does not achieve maximal reward, even when it receives rewards during exploration.
7. Related Work
We draw on the long line of work on learning to adapt to new similar tasks (Schmidhuber, 1987; Thrun & Pratt, 2012; Naik & Mammone, 1992; Bengio et al., 1991; 1992;
Hochreiter et al., 2001; Andrychowicz et al., 2016; Santoro et al., 2016).
Within meta-RL, many works focus on adapting efﬁciently to a new task from few samples without optimizing the sample collection process, via updating the policy parameters (Finn et al., 2017; Rothfuss et al., 2018;
Agarwal et al., 2019; Yang et al., 2019; Mendonca et al., 2019), learning a model (Nagabandi et al., 2018; Sæmundsson et al., 2018), multi-task learning (Fakoor et al., 2019), or leveraging demonstrations (Zhou et al., 2019a).
Instead, we focus on problems where targeted exploration is critical for few-shot adaptation.
More closely related to our work are approaches that speciﬁcally explore to obtain the most informative samples.
These [Figure] fall into two main categories: end-to-end and decoupled approaches.
The ﬁrst category jointly learns exploration and adaptation behaviors end-to-end by training a recurrent policy to maximize returns (Duan et al., 2016;
Wang et al., 2016a; Mishra et al., 2017; Stadie et al., 2018; Zintgraf et al., 2019; Humplik et al., 2019; Kamienny et al., 2020; Ritter et al., 2018).
and can represent the optimal policy (Kaelbling et al., 1998), but, as this work shows, suffer from coupling between exploration and adaptation.
Notably, many of these works use separate objectives to train the representation of explored transitions but not for directly improving the exploration behavior.
As a result, the coupling problem remains.
The second category decouples exploration from adaptation via, e.g., Thompson-sampling (TS) (Thompson, 1933; Rakelly et al., 2019), obtaining exploration trajectories predictive of dynamics or rewards (Zhou et al., 2019b;
Gurumurthy et al., 2019), or exploration noise (Gupta et al., 2018).
These avoid the problem of coupling in the ﬁrst category, but no longer learn optimal exploration.
In particular, TS (Rakelly et al., 2019) explores by guessing the task and executing a policy for that task, and therefore cannot represent exploration behaviors that are different from adaptation (Russo et al., 2017).
Predicting the dynamics (Zhou et al., 2019b; Gurumurthy et al., 2019) is inefﬁcient when only a small subset of the dynamics are relevant to solving the task.
In contrast, D REAM ’s objective is consistent and yields optimal exploration when maximized.
Past work (Gregor et al., 2016; Houthooft et al., 2016;
Eysenbach et al., 2018; Warde-Farley et al., 2018) also optimizes mutual information objectives, but not for meta-RL.
IMRL also draws inspiration from prior work.
Rewardfree adaptation relates to prior work that considers ignoring the rewards in the ﬁrst episodes (Stadie et al., 2018) and when rewards are unavailable at test time (Yang et al., 2019).
Other work (Nagabandi et al., 2018; Sæmundsson et al., 2018) assume that the reward function is known to the agent, like instructions.
8. Conclusion
In summary, we present two contributions.
First, we identify a coupling issue with existing meta-RL works that leads to poor sample complexity, and propose decoupled objectives for exploration and execution (D REAM ) to address this.
Unlike prior decoupled objectives (Zhou et al., 2019b; Rakelly et al., 2019; Gupta et al., 2018; Gurumurthy et al., 2019) that do not learn the optimal exploration strategy, maximizing our decoupled objectives also maximizes returns.
Second, we propose a new meta-RL setting (IMRL) to model realistic situations where the agent receives instructions.
The combination of these two contributions is greater than their sum: In IMRL, D REAM adapts to new environments and tasks without ever seeing rewards during meta-testing, solving challenging problems that prior methods fail to solve.
IMRL also alleviates unnecessary exploration for inferring the task.
We therefore advocate for research on adaptation to provide the agent instructions like in IMRL.
While our work provides many beneﬁts, we study the setting where the agent is allowed initial training episodes.
Other work like V ARI BAD can, in principle, adapt even without these initial episodes, which we leave to future work.
Future work could also explore more sophisticated representations for the problem IDs and instructions, such as natural language. Reproducibility.
Our code is available at https:// github.com/ezliu/dream .
Acknowledgements
This material is based upon work supported by the National
Science Foundation Graduate Research Fellowship under
Grant No. 2018255189.
Icons used in this paper were made by Freepik, ThoseIcons and dDara from www.flaticon.
com .
A. D REAM Training Details
Algorithm 1 summarizes a practical algorithm for training
D REAM , parametrizing the policies as deep dueling doubleQ networks (Wang et al., 2016b; van Hasselt et al., 2016), with exploration Q-values ˆ Q exp ( s, τ exp , a ; φ ) parametrized by φ (and target network parameters φ (cid:48) ) and execution Qvalues ˆ Q exe ( s, i, z, a ; θ ) parametrized by θ (and target network parameters θ (cid:48) ).
We train on trials with one exploration and one execution episode, but can test on arbitrarily many execution episodes, as the execution policy acts on each episode independently (i.e. it does not maintain a hidden state across episodes).
Using the choices for F ψ and q ω in
Section 4.3, training proceeds as follows.
We ﬁrst sample a new problem for the trial and roll-out the exploration policy, adding the roll-out to a replay buffer (lines 7-9).
Then, we sample an instruction and roll-out the execution policy, adding the roll-out to a replay buffer (lines 10-13).
We train the execution policy on both stochastic encodings of the problem ID N ( f ψ ( T ) , ρ 2 I ) and on encodings of the exploration trajectory g ω ( τ exp ) .
Next, we sample from the replay buffers and update the parameters.
First, we sample from the exploration replay buffer and perform a normal DDQN update on the exploration Q-value parameters φ using rewards computed from the decoder (lines 14-16).
Concretely, we take the following gradient step, regressing toward the target network:
[Equation]
where target =
r exp t + γ ˆ Q exp ( s t [Equation] 1 , [ τ exp : t − 1 ; a t ; s t ] , a DDQN ; φ (cid:48) ) , and a DDQN = arg max a ˆ Q exp ( s t +1 , [ τ exp : t − 1 ; a t ; s t ]; φ ) .
We perform a similar update with the execution Q-value parameters (lines 17-19).
We sample from the execution replay buffer and perform two DDQN updates, one from the encodings of the exploration trajectory and one from the encodings of the problem ID.
These take the following form:
[Equation]
∇ [Equation] [Equation] , i,
[Equation]
ψ [Equation] ) , a ; θ ) − target prob (cid:13)(cid:13)(cid:13) 2 2 , where target traj = r + ˆ Q exe ( s (cid:48) , i, g ω (cid:48) ( τ exp ) , a traj ; θ (cid:48) )
a traj = arg max a ˆ Q exe ( s (cid:48) , i, g ω ( τ exp ) , a ; θ ) , and target prob =
r + ˆ Q exe ( s (cid:48) , i, f ψ (cid:48)
( T ) , a prob ; θ (cid:48) )
a prob = arg max a ˆ Q exe
( s (cid:48) , i, f ψ ( T ) , a ; θ ) .
Finally, from the same execution replay buffer samples, we also update the problem ID embedder to enforce the information bottleneck (line 20) and the decoder to approximate the true conditional distribution (line 21).
Since the magnitude (cid:107) f ψ ( T ) (cid:107) 22 partially determines the scale of the reward, we add a hyperparameter K and only minimize the magnitude when it is larger than K .
As is usual with deep
Q-learning, instead of sampling from the replay buffers and updating after each episode, we sample and perform all of these updates every 4 timesteps.
We periodically update the target networks (lines 22-23).
B. Experiment Details B.1. Tasks
Distracting bus / map.
Riding each of the four colored buses teleports the agent to one of the green goal locations.
In different problems, the destinations of the colored buses change.
Additionally, in the distracting bus domain, the problem also encodes the x-coordinates of the destinations of the gray buses in the bottom gray row.
Overall, in the map domain, the problem T is a one-hot representation representing which of the 4!
permutations of the four green goal locations the colored buses map to.
The states include an extra dimension, which is set to 0 when the agent is not at the map, and is set to this one-hot value T when the agent is at the map.
In the distracting bus domain, the problem is represented as a 5-tuple, where the ﬁrst index is the same as the map domain, and the remaining four indices describe the x-coordinates of the gray buses.
Figure 7 displays three such examples.
This ﬁgure also shows the optimal exploration policies learned by D REAM .
When we remove the information bottleneck from D REAM , as in Figure 5, the exploration begins to visit the gray buses, because these are part of the problem representation, even though this results in lower execution returns.
Cooking.
In different problems, the (color-coded) fridges contain 1 of 7 different ingredients.
The ingredients in each fridge are unknown until the goes to the fridge and uses the pickup action.
Figure 8 displays three example problems and the optimal exploration learned by D REAM .
The representation of the problem T is a triple, where the i -th index is an indicator of what ingredient is in the i -th fridge.
Each instruction i corresponds to a recipe of placing the two correct ingredients in the pot in the right order.
We represent this as a tuple, where the ﬁrst index is an indicator of the ﬁrst ingredient and the second index is the indicator of the second ingredient.
In a given problem, we only sample recipes (instructions) involving the ingredients actually present in that problem.
For test time, we hold out a randomly chosen 1% of the 7 3 different problems and 3 of the 7 2 different instructions: (3 , 6) , (7 , 1) , and (5 , 4) .
: t ) (cid:107) 22 − (cid:13)(cid:13)
f ψ ( T ) − g ω ( τ exp : t − 1 ) (cid:13)(cid:13) 2 2 − c (Equation 5). Optimize φ with DDQN update with tuple ( s t , a t , r exp t , s t +1 ) 17: Sample batches of ( s, a, r, s (cid:48) , i, T , τ exp ) ∼ B exe from execution replay buffer.
Optimize θ and ω with DDQN update with tuple (( s, i, τ exp ) , a, r, ( s (cid:48) , i, τ exp )) 19:
Optimize θ and ψ with DDQN update with tuple (( s, i, T ) , a, r, ( s (cid:48) , i, T )) 20:
Optimize ψ on ∇ ψ min( (cid:107) f ψ ( T ) (cid:107) 22 , K ) (Equation 3) 21:
{+Optimize ω on ∇ ω (cid:80) t (cid:107) f ψ ( T ) − g ω ( τ exp : t ) (cid:107) 
+}
(Equation 4) 22: if trial ≡ 0 (mod target freq ) then 23:
Update target parameters φ (cid:48)
=
φ , θ (cid:48)
= θ , ψ (cid:48)
= ψ , ω (cid:48)
=
{+ω 24: end if 25: end for    
+}
[Figure]
[Figure]
[Figure]
[Figure]
We use a shaped reward function A T .
The agent receives a per timestep penalty of − 0 .
1 reward and receives +0 .reward for completing each of the four steps: (i) picking up the ﬁrst ingredient speciﬁed by the instruction; (ii) placing the ﬁrst ingredient in the pot; (iii) picking up the second ingredient speciﬁed by the instruction; and (iv) placing the second ingredient in the pot.
To prevent the agent from gaming the reward function, e.g., by repeatedly picking up the ﬁrst ingredient, dropping the ﬁrst ingredient anywhere but in the pot yields a penalty of − 0 .
25 reward, and similarly for all steps.
Sparse-reward 3D visual navigation.
We implement this domain in Gym MiniWorld (Chevalier-Boisvert, 2018), where the agent’s observations are 80 × 60 × 3 RGB arrays.
There are two problems T =
(the sign says “blue”) and
T = 1 (the sign says “red”).
There are three instructions i = 0 , 1 , 2 corrsponding to picking up the ball, key, and box, respectively.
The reward function A T ( s, i ) is +1 for picking up the correct colored object (according to T ) and the correct type of object (according to the instruction) and − 1 for picking up an object of the incorrect color or type.
Otherwise, the reward is 0.
On each episode, the agent begins at a random location on the other side of the barrier from the sign.
B.2. Additional Results
Distracting bus / map.
Figures 7a, 7b, and 7c show the exploration policy D REAM learns on the distracting bus and map domains.
With the information bottleneck, D REAM optimally explores by riding 3 of the 4 colored buses and inferring the destination of the last colored bus (Figure 7a).
Without the information bottleneck, D REAM explores the unhelpful gray buses and runs out of time to explore all of the colored buses, leading to lower reward (Figure 7b).
In the map domain, D REAM optimally explores by visiting the map and terminating the exploration episode.
In contrast, the other methods (RL 2 , I MPORT , V ARI BAD) rarely visit the colored buses or map during exploration and consequently walk to their destination during execution, which requires more timesteps and therefore receives lower returns.
In Figure 9, we additionally visualize the exploration trajectory encodings g ω ( τ exp ) and problem ID encodings f ψ ( T ) that D REAM learns in the distracting bus domain by applying t-SNE (van der Maaten & Hinton, 2008).
We visualize the encodings of all possible problem IDs as dots.
They naturally cluster into 4!
= 24 clusters, where the problems within each cluster differ only in the destinations of the gray distracting buses, and not the colored buses.
Problems in the support of the true posterior p ( T | τ exp ) are drawn in green, while problems outside the support (e.g., a problem that speciﬁes that riding the green bus goes to location (3 , 3) when it has already been observed in τ exp that riding the orange bus goes to location (3 , 3) ) are drawn in red.
[-Optimize ω on ∇ ω (cid:80) t (cid:107) f ψ ( T ) − g ω ( τ exp : t ) (cid:107) 
-]
[-ω 24: end if 25: end for    
-]
We also plot the encoding of the exploration trajectory τ exp so far as a blue cross and the mean of the green clusters as a black square.
We ﬁnd that the encoding of the exploration trajectory g ω ( τ exp ) tracks the mean of the green clusters until the end of the exploration episode, when only one cluster remains, and the destinations of all the colored buses has been discovered.
Intuitively, this captures uncertainty in what the potential problem ID may be.
More precisely, when the decoder is a Gaussian, placing g ω ( τ exp ) at the center of the encodings of problems in the support exactly minimizes Equation 4.
Cooking.
Figure 8a shows the exploration policy D REAM learns on the cooking domain, which visits each of the fridges and investigates the contents with the ”pickup” action.
In contrast, the other methods rarely visit the fridges during exploration, and instead determine the locations of the ingredients during execution, which requires more timesteps and therefore receives lower returns.
B.3. Model Architecture
In this section, we describe the model architectures used in our experiments.
Where possible, we use the same model architecture for all methods: D REAM , RL 2 , I MPORT , and
V ARI BAD.
D REAM
For the decoder g ω ( τ exp =
( s 0 , a 0 , s 1 , . .
, s T )) , we embed each transition ( s t , a t , s t +1 ) of the exploration trajectory τ exp as follows.
We apply a state learned embedding function to both states e ω ( s t ) and e ω ( s t +1 ) and embed the action e ω ( a t ) by applying a learned embedding matrix with output dimension 32.
Then we apply two linear layers with output dimension 128 and 64 respectively, using
ReLU activations to obtain the embedding for each transition.
Given embeddings for each transition, we embed the entire trajectory by passing an LSTM with output dimension 128 on top of the transition embeddings, followed by two [Figure] linear layers of output dimension 128 and 64 with ReLU activations.
For the execution policy Q-values ˆ Q exe θ ( a | s, i, z ) , we either choose z to be the decoder embedding of the exploration trajectory g ω ( τ exp ) of to be an embedding of the problem
ID e θ ( T ) , where we always use the exploration trajectory embedding g ω ( τ exp ) at meta-test time.
We embed the state and instructions with learned embedding functions e ( s ) and e θ ( i ) respectively.
Then we apply a linear layer of output dimension 64 to the concatenation of [ e ( s ); e θ ( i ); z ] with a
ReLU activation.
Finally, we apply two linear layer heads of output dimension 1 and |A| respectively to form estimates of the value and advantage functions, using the dueling Qnetwork parametrization.
To obtain Q-values, we add the value function to the advantage function, subtracting the mean of the advantages.
For the exploration policy Q-values ˆ Q exp φ ( a t | s t , τ exp : t ) , we embed the past states, actions, and rewards recurrently with an LSTM of output dimension 64.
We embed each state, action, and reward with learned embedding functions e φ ( s ) , e φ ( a ) , e φ ( r ) respectively, where the actions are embedded with a learned embedding matrix of output dimension 16 and the rewards are embedded with a linear layer with output dimension 16 .
Then, we pass a linear layer with output dimension 64 and ReLU activations over their concatenation [ e φ ( s t ); e φ ( a t ); e φ ( r t )] at each timestep t .
We then pass an
LSTM with output dimension 64 over this, setting the initial hidden state to all 0 .
We apply two linear layer heads to obtain value and advantage estimates as above.
RL
RL 2 learns a policy π ( a t | s t , i, τ : t ) producing actions a t given the state s t , instructions i and history τ : t .
We parametrize this dueling double Q-networks, learning Qvalues ˆ Q ( s t , i, τ : t , a t ) .
We use an identical architecture to
D REAM ’s exploration policy Q-values, except we also embed the instructions and episode terminations.
Concretely, we learn embedding functions on the state e ( s ) , action e ( a ) , rewards e ( r ) , episode terminations e ( d ) , and instructions e ( i ) and pass a linear layer of output dimension 64 over the concatenation of [ e ( s t ); e ( a t ); e ( r t ); e ( d t );
e ( i t )] with a
ReLU activation at each timestep t .
Then, we pass an LSTM with output dimension 64 over this.
We embed actions and episode terminations with a learned embedding matrix with output dimension 16 and embed rewards with a learned linear layer with output dimension 16 .
As for all policies, we apply two ﬁnal linear layer heads to obtain value and advantage estimates.
I MPORT
I MPORT learns a policy π ( a t | s t , i, z ) producing actions a t given the state s t , instructions i and embedding z of either the problem T or the history τ : t We also parametrize this dueling double Q-networks, learning Qvalues ˆ Q ( s t , i, z, a t ) , using a similar architecture to RL 2 .
We embed the history τ : t and instruction i to obtain e ( τ : t ) identically as RL 2 above.
We embed the problem T and current state and instructions ( s t , i ) with learned embedding functions e ( T ) , e ( s t ) , and e ( i ) respectively.
Then we alternate meta-training trials between choosing z
= e ( T ) and z = e ( τ : t ) .
We apply a linear layer of output dimensionto the concatenation [ e ( s t ); e ( i ); z ] with ReLU activations and then apply two linear layer heads to obtain value and advantage estimates.
V ARI BAD
V ARI BAD also learns a policy π ( a t | s t , i, τ : t ) producing actions a t given the state s t , instructions i and history τ : t .
We also parametrize this dueling double Q-networks, learning Q-values ˆ Q ( s t , i, z, a t ) , using a similar architecture to above.
To embed the history and instructions, V ARI BAD learns an encoder enc ( z | τ : t ) , which we parametrize by embedding the history τ : t using the same architecture as RL 2 to obtain e ( τ : t ) .
Then we apply a linear layer with output dimension to obtain a mean µ .
Finally, we sample z from N ( µ, ν 2 I ) , where ν 2 = 0 .
We also tried learning the variance instead of ﬁxing it to ν 2 I by applying a second linear head to the e ( τ : t ) , but found no change in performance, so stuck with the simpler ﬁxed variance approach.
Given z sampled from the encoder, we embed the current state s t with a learned encoder e ( s t ) and apply a linear layer of output dimension 64 to the concatenation [ e ( s t ); e ( i ); z ] .
Then, we apply two linear layer heads to obtain value and advantage estimates.
V ARI BAD also learns a state decoder ˆ T ( s (cid:48) | a, s, z ) and reward decoder ˆ R ( s (cid:48) | a, s, z ) for its auxiliary losses, where z is sampled from the decoder.
For both of these, we embed the action e ( a ) with an embedding matrix of output dimension 32 and learn a state embedder e ( s ) .
Then we apply two linear layers with output dimension 128 to the concatenation [ e ( s ); e ( a ); z ] .
Finally, we apply two linear heads, one for the state decoder and one for the reward decoder and take the mean-squared error.
In the 3D visual navigation domain, we remove the state decoder, because the state is too high-dimensional to predict.
Note that Zintgraf et al.
(2019) found better results when removing the state decoder in all experiments.
We also tried to remove the state decoder in the grid world experiments, but found better performance when keeping the state decoder.
We also found that V ARI BAD performed better without the KL loss term, so we excluded that for our ﬁnal experiments.
State, instruction, and problem ID embeddings.
For all learned embeddings of the state, instruction and problem
ID embeddings, we embed each dimension independently with an embedding matrix of dimension 32.
Then we concatenate the per-dimension embeddings and apply two linear layers with output dimension 256 and 64 with ReLU activations.
In the 3D visual navigation task, we use the architecture recommended by Chevalier-Boisvert (2018) for the learned state embeddings instead of the embedding scheme described above.
We apply 3 CNN layers each with 32 output layers and stride length 2 and with kernel sizes of 5, 5, and 4 respectively.
We apply ReLU activations in between the 3 CNN layers and apply a ﬁnal linear layer to the ﬂattened output of output dimension 128.
B.4. Hyperparameters
In this section, we detail the hyperparameters used in our experiments.
Where possible, we used the default DQN hyperparameter values from Mnih et al. (2015).
and shared the same hyperparameter values across all methods for fairness.
We optimize all methods with the Adam optimizer (Kingma & Ba, 2014).
Table 1 summarizes the shared hyperparameters used by all methods and we detail any differences in
[Table]
All methods use a linear decaying (cid:15) schedule for (cid:15) -greedy exploration.
For RL 2 , I MPORT , and V ARI BAD, we decay (cid:15) from 1 to 0 .
01 over 500000 steps.
For D REAM , we split the decaying between the exploration and execution policies.
We decay each policy’s (cid:15) from 1 to 0.01 over 250000 steps.
We train the recurrent policies (D REAM ’s exploration policy,
RL 2 , I MPORT , and V ARI BAD) with a simpliﬁed version of the methods in Kapturowski et al.
(2019) by storing a replay buffer with up to 16000 sequences of 50 consecutive timesteps.
We decrease the maximum size from 16000 to 10000 for the 3D visual navigation experiments in order to ﬁt inside a single NVIDIA GeForce RTX 2080 GPU.
For
D REAM ’s execution policy, the replay buffer stores up to 100K experiences (60K for 3D visual navigation).
For D REAM , we additionally use per timestep exploration reward penalty c = 0 .
01 , decoder and stochastic encoder variance ρ 2 = 0 .
1 , and information bottleneck weight λ = 1 .
C. Analysis C.1. Consistency
We restate the consistency result of D REAM (Section 5.1) and prove it below.
Proposition 1.
Let V ∗
( i ; T ) be the maximum expected returns achievable by any execution policy with access to T on instruction i and problem T , i.e., with complete information.
Let π exe (cid:63) , π exp (cid:63) , F (cid:63) and q (cid:63) ( z | τ exp ) be the optimizers of the D REAM objective.
Then for large enough T (the length of the exploration episode) and M (number of exploration episodes), and expressive enough function classes,
[Equation]
Proof.
Recall that π exe (cid:63) and F (cid:63) ( z | T ) are optimized with an information bottleneck according to Equation 2.
Note that if π exe (cid:63) is optimized over an expressive enough function class and λ approaches 0 , then π exe (cid:63) achieves the desired expected returns conditioned on the stochastic encoding of the problem F (cid:63) ( z | T ) (i.e., has complete information):
[Equation]
where V π exe (cid:63) ( i, z ; T ) is the expected returns of π exe (cid:63) on problem T given instruction i and embedding z .
Therefore, it sufﬁces to show that the distribution over z from the decoder q (cid:63) ( z | τ exp ) is equal to the distribution from the encoder
F (cid:63) ( z | T ) for all exploration trajectories in the support of π exp ( τ exp | T ) 1 , for each problem T .
Then,
[Equation]
as desired.
We show that this occurs as follows.
Given stochastic encoder F (cid:63) ( z | T ) , exploration policy π exp (cid:63) maximizes I ( τ exp ; z )
=
H ( z ) − H
( z | τ exp )
(Equation 4) by assumption.
Since only H ( z | τ exp ) depends on π exp (cid:63) , the exploration policy outputs trajectories that minimize
H ( z | τ exp ) =
E T ∼ p ( T ) (cid:2) E τ exp ∼ π exp ( τ exp ∼T ) (cid:2) E z ∼ F (cid:63) ( z |T ) [ − log p ( z | τ exp )] (cid:3)(cid:3) = E T ∼ p ( T ) (cid:2) E τ exp ∼ π exp ( τ exp ∼T )
[ H ( F (cid:63) ( z | T ) , p ( z | τ exp ))] (cid:3) , where p ( z | τ exp ) is the true conditional distribution and
H ( F (cid:63) ( z | T ) , p ( z | τ exp )) is the cross-entropy.
The crossentropy is minimized when p ( z | τ exp ) =
F (cid:63) ( z | T ) , which can be achieved with enough exploration episodes M and long enough exploration trajectories T (by visiting each transition sufﬁciently many times).
Optimized over an expressive enough function class, q (cid:63) ( z | τ exp ) equals the true conditional distribution p ( z | τ exp ) at the optimum of Equation 4, which equals F (cid:63) ( z | T ) as desired.
C.2. Tabular Example
We ﬁrst formally detail a more general form of the simple tabular example in Section 5.2, where episodes are horizon H rather than 1-step bandit problems.
Then we prove sample complexity bounds for RL 2 and D REAM , with (cid:15) - greedy tabular Q-learning with (cid:15)
= 1 , i.e., uniform random exploration.
Setting.
We construct this horizon H setting so that taking a sequence of H actions a (cid:63) (instead of a single action as on problem T . before) in the exploration episode leads to a trajectory τ exp (cid:63) that reveals the problem T ; all other action sequences a lead to a trajectory τ exp a that reveals no information.
Similarly, the problem T identiﬁes a unique sequence of H actions a T that receives reward 1 during execution, while all other action sequences receive reward 0 .
Again, taking the action sequence a (cid:63) during exploration is therefore necessary and sufﬁcient to obtain optimal reward 1 during execution.
We formally construct this setting by considering a family of episodic MDPs (cid:104)S , A , R T , T T (cid:105) parametrized by the problem ID
T ∈ M , where:
• Each execution and exploration episode is horizon H .
• The action space
A consists of A discrete actions { 1 , 2 , . . .
, A } .
• The space of problems M = { 1 , 2 , . .
, |A| H } and the distribution p ( T ) is uniform.
To reveal the problem via the optimal action sequence a (cid:63) and to allow a T to uniquely receive reward, we construct the state space and deterministic dynamics as follows
• States s ∈ S are ( H + 2) -dimensional
and the deterministic dynamics are constructed so the ﬁrst index represents the current timestep t , the middle H indices represent the history of actions taken, and the last index reveals the problem ID if a (cid:63) is taken.
The initial state is the zero vector s
=
and we denote the state at the t -th timestep s t as ( t, a 0 , a 1 , . .
. , a t − 1 , 0 , . .
• The optimal exploration action sequence a (cid:63) is set to be taking action 1 for H timesteps.
In problem T taking action a H − 1 = 1 at state s H − 1 = ( H − 1 , a 0 = 1 , . . .
, a H − 2 = 1 , 0 , 0) (i.e., taking the entire action sequence a (cid:63) ) transitions to the state s H =
( H, a 0 = 1 , . . .
, a H − 2 = 1 , a H − 1 = 1 , T ) , revealing the problem T .
• The action sequence a T identiﬁed by the problem T is set as the problem T in base |A| : i.e., a T is a sequence of H actions ( a 0 , a 1 , . . .
, a H − 1 ) with (cid:80) H −t =0 a t |A| t = T .
In problem T with a T = ( a 0 , a 1 , . . . , a H − 1 ) , taking action a H − 1 at timestep H − 1 at state s H − 1 = ( H − 1 , a 0 , a 1 , . . .
, a H − 2 , 0 , 0) (i.e., taking the entire action sequence a T ) yields R T
( s H − 1 , a H − 1 )
= 1 .
Reward is 0 everywhere else: i.e., R T ( s, a ) = 0 for all other states s and actions a .
• With these dynamics, the exploration trajectory τ exp a =
( s 0 , a 0 , r 0 , . . .
, s H ) is uniquely identiﬁed by the action sequence a
and the problem T if revealed in s H .
We therefore write τ exp a
=
( a , T ) for when a = a (cid:63) reveals the problem T , and τ exp a
=
( a , 0) , otherwise.
Uniform random exploration.
In this general setting, we analyze the number of samples required to learn the optimal exploration policy with RL 2 and D REAM via (cid:15) -greedy tabular Q-learning.
We formally analyze the simpler case where (cid:15) = 1 , i.e., uniform random exploration, but empirically ﬁnd that D REAM only learns faster with smaller (cid:15) , and RLonly learns slower.
In this particular tabular example with deterministic dynamics that encode the entire action history and rewards, learning a per timestep Q-value is equivalent to learning a
Q-value for the entire trajectory.
Hence, we denote exploration Q-values ˆ Q exp ( a ) estimating the returns from taking the entire sequence of H actions a at the initial state sand exeuction Q-values ˆ Q exe ( τ exp , a ) estimating the returns from taking the entire sequence of H actions a at the initial state s 0 given the exploration trajectory τ exp .
We drop sfrom notation, since it is ﬁxed.
Recall that RL 2 learns exploration Q-values ˆ Q exp by regressing toward the execution Q-values ˆ Q exe .
We estimate the execution Q-values ˆ Q exe ( τ exp , a ) as the sample mean of returns from taking actions a given the exploration trajectory τ exp and estimate the exploration Q-values ˆ
Q exp ( a ) as the sample mean of the targets.
More precisely, for action sequences a (cid:54) = a (cid:63) , the resulting exploration trajectory τ exp a is deterministically ( a , 0) , so we set ˆ
Q exp ( a ) =
ˆ V exe ( τ exp a ) =
max a (cid:48)
ˆ Q exe ( τ exp a , a (cid:48) ) .
For a (cid:63) , the resulting exploration trajectory τ exp a (cid:63) may be any of ( a (cid:63) , T ) for T ∈ M , so we set ˆ Q exp ( a (cid:63) ) as the empirical mean of ˆ
V exe ( τ exp a (cid:63) ) of observed τ exp a (cid:63) .
Recall that D REAM learns exploration Q-values ˆ Q exp by regressing toward the learned decoder log ˆ q ( T | τ exp a ) .
We estimate the decoder ˆ q ( T | τ exp a ) as the empirical counts of ( T , τ exp a ) divided by the empirical counts of τ exp a and similarly estimate the Q-values as the empirical mean of log
ˆ q ( T | τ exp a ) .
We denote the exploration Q-values learned after t timesteps as ˆ Q exp t , and similarly denote the estimated decoder after t timesteps as ˆ q t .
Given this setup, we are ready to state the formal sample complexity results.
Intuitively, learning the execution Qvalues for RL 2 is slow, because, in problem T , it involves observing the optimal exploration trajectory from taking actions a (cid:63) and then observing the corresponding execution actions a T , which only jointly happens roughly once per |A| 2 H samples.
Since RL 2 regresses the exploration
Q-values toward the execution Q-values, the exploration
Q-values are also slow to learn.
In contrast, learning the decoder ˆ q ( T | τ exp a ) is much faster, as it is independent of the execution actions, and in particular, already learns the correct value from a single sample of a (cid:63) .
We formalize this intuition in the following proposition, which shows that D REAM learns in a factor of at least |A| H |M| fewer samples than RL 2 .
Proposition 1.
Let T be the number of samples from uniform random exploration such that the greedyexploration policy is guaranteed to be optimal (i.e., arg max a ˆ Q exp t ( a )
= a (cid:63) ) for all t ≥ T .
If ˆ Q exp is learned with D REAM , the expected value of T is O ( |A| H log |A| H ) .
If ˆ Q exp is learned with RL 2 , the expected value of T is Ω( |A| 2 H |M| log |A| H ) .
Proof.
For D REAM , ˆ Q exp T ( a (cid:63) ) >
ˆ Q exp T ( a ) for all a (cid:54) = a (cid:63) if log ˆ q T ( T | ( a (cid:63) , T )) > log ˆ q T ( T | ( a , 0)) for all T and a (cid:54) = a (cid:63) .
For all t ≥ T , ˆ Q exp t is guaranteed to be optimal, since no sequence of samples will cause log ˆ q t ( T |
( a (cid:63) , T )) = 0 ≤ log ˆ q t ( T | ( a , 0)) for any a (cid:54) = a (cid:63) .
This occurs once we’ve observed ( T , ( a , 0)) for two distinct T for each a (cid:54) = a (cid:63) and we’ve observed ( T , ( a (cid:63) , T )) for at least one T .
We can compute an upperbound on the expected number of samples required to observe ( T , τ exp a ) for two distinct T for each action sequence a by casting this as a coupon collector problem, where each pair ( T , τ exp a ) is a coupon.
There are 2 |A| H total coupons to collect.
This yields that the expected number of samples is O ( |A| H log |A| H ) .
For RL 2 , the exploration policy is optimal for all timesteps t ≥ T for some T only if the instruction values ˆ V exe T ( τ exp =
( a (cid:63) , T ))
= 1 for all T in M .
Otherwise, there is a small, but non-zero probability that ˆ V exe t ( τ exp =
( a , 0)) will be greater at some t > T .
For the instruction values to be optimal at all optimal exploration trajectories ˆ V exe T ( τ exp =
( a (cid:63) , T ))
= 1 for all T ∈ M , we must jointly observe exploration trajectory τ exp =
( a (cid:63) , T ) and corresponding action sequence a T for each problem T ∈ M .
We can lowerbound the expecected number of samples required to observe this by casting this as a coupon collector problem, where each pair ( τ exp =
( a (cid:63) , T ) , a T ) is a coupon.
There are |M| · |A| H unique coupons to collect and collecting any coupon only occurs with probability 1 |A| H in each episode.
This yields that the expected number of samples is Ω( |A| 2 H · |M| · log |A| H ) .