Meta-Reinforcement Learning With Informed Policy Regularization
Authors: Pierre-Alexandre Kamienny
Matteo Pirotta
Alessandro Lazaric
Thibault Lavril
Nicolas Usunier
Ludovic Denoyer
Venue:  (2021)
Keywords: meta-reinforcement learning
reinforcement learning
multi-task
non-stationary
task representations
regularization

Abstract: Meta-reinforcement learning aims at finding a policy able to generalize to new environments. When facing a new environment, this policy must explore to identify its particular characteristics and then exploit this information for collecting reward.  We consider the online adaptation setting where the agent needs to trade-off between the two types of behaviour within the same episode. Even though policies based on recurrent neural networks can be used in this setting by training them on multiple environments, they often fail to model this trade-off, or solve it at a very high computational cost. In this paper, we propose a new algorithm that uses privileged information in the form of a task descriptor at train time to improve the learning of recurrent policies.  Our method learns an informed policy (i.e., a policy receiving as input the description of the current task) that is used to both construct task embeddings from the descriptors, and to regularize the training of the recurrent policy through parameters sharing and an auxiliary objective. This approach significantly reduces the learning sample complexity without altering the representational power of RNNs, by focusing on the relevant characteristics of the task, and by exploiting them efficiently.  We evaluate our algorithm in a variety of environments that require sophisticated exploration/exploitation strategies and show that it outperforms vanilla RNNs, Thompson sampling and the task-inference approaches to meta-reinforcement learning.

Under review as a conference paper at ICLR 2021
META-REINFORCEMENE LEARNING WITH INFORMED POLICY REGULARIZATION
Anonymous authors Paper under double-blind review
ABSTRACT
Meta-reinforcemmntt learning aims at finding a policy able to generalize to new environments. When facing new environment, this policy must explore to identify its particular characteristics and then exploit this information for collecting reward. We consider the online adaptation setting where the agent needs to trade-off' between the two types of behaviour within the same episode. Even though policies based on recurrent neural networks can be used in this setting by training them multiple environments, they often fail to model this trade-off, very high computational cost. In this paper, we propose a new algorithm that uses privileged information in the form of a task descriptor train time to improve the learning of recurrent policies. Our method learns an informed policy (i.e., a policy receiving input the description of the current task) that is used to both construct task embeddings from the descriptors, and to regularize the training of the recurrent policy through parameters sharing and an auxiliary objective. This approach significantly reduces the learning sample complexity without altering the representational power of RNNs, by focusing on the relevant characteristics of the task, and by exploiting them efficiently. We evaluate our algorithm in a variety of environments that require sophisticated exploration/exploiatio strategies and show that it outperforms vanilla RNNs, Thompson sampling and the task-inference approaches to meta-reinforcement learning.
solve it at
INTRODUCTION
Deep Reinforcement Learning has been used to successfully train agents on a range of challenging environments such as Atari games (Mnih et al., 2013: Bellemare et al.) 2013 Hessel et al.) (017) or continuous control (Peng et al. 2017] Schulman et al. 2017). Nonetheless, in these problems, RL agents perform exploration strategies to discover the environment and implement algorithms to learn a policy that is tailored to solving a single task. Whenever the task changes, RL agents generalize poorly and the whole process of exploration and learning restarts from scratch. On the other hand, we expect an intelligent agent to fully master a problem when it is able to generalize from few instances (tasks) and achieve the objective of the problem under many variations of the environment. For instance, children know how to ride bike (i.c.. the problem) when they can reach their destination irrespective of the specific bike they are riding, which requires to adapt to the weight of the bike, the friction of the brakes and tires, and the road conditions (i.c., the tasks).
How to enable agents generalize across tasks has been studied in Multi-task Reinforcement Learning (e.g. Wilson et al., 2007 Teh et al., 2017), Transfer Learning (e.g. Taylor & Stone 2011 Lazaric 2012) and Meta Reinforcement Learning (Finn al.] [017 Hausman et al. 20188]
Rakelly et al. 2019 Humplik et al., 2019). These works fall into two categories. Learning to learn approaches aim at speeding up learning on new tasks, by pre-training feature extractors or learning good initializations of policy weights (Raghu al. 2019).
In contrast, we study in this paper the online adaptation setting where a single policy is trained for a fixed family of tasks. When facing a new task, the policy must then balance exploration (or probing), to reduce the uncertainty about the current task, and exploitation to maximize the cumulative reward of the task. Agents evaluated their ability manage this trade-off within single episode of the same task. The online adaptation setting is a special case of a partially observed markov decision problem, where the unobserved variables are the descriptors of the current task. It is thus

---

Under review as a conference paper at ICLR 2021
G2
Figure 1: An environment with two tasks: the goal location (G1 or G2) changes at each episode. The sign reveals the location of the goal. Optimal informed policies are shortest paths from start either G1 or G2, which never visit the sign. Thompson sampling cannot represent the optimal exploration/explotaiieee
policy (go to the sign first) since going to the sign is not feasible by any informed policy.
igo
possible to rely recurrent neural networks (RNNs) Bakker 2001 Heess et al. 2013). since they can theoretically represent optimal policies in POMDPs if given enough capacity. Unfortunately. the training of RNN policies has often prohibitive sample complexity and it may converge to suboptimal local minima.
To overcome this drawback, efficient online adaptation methods leverage the knowledge of the task at training time. The main approach is to pair an exploration strategy with the training of informed policies, i.e. policies taking the description of the current task input. Probe-then-Exploit (PTE) algorithms (e.g. Zhou et al. 2019) operate in two stages. They first rely on an exploration policy to identify the task. Then, they commit to the identified task by playing the associated informed policy. Thompson Sampling (TS) approaches [Thompson] 1933, Osband et al. 2016 2019) maintain a distribution over plausible tasks and play the informed policy task sampled from the posterior following predefined schedule. PTE and TS are expected to be sample-efficient relatively to RNNs as learning informed policies is a fully observable problem. However, as we discuss in Section 3 PTE and TS cannot represent effective exploration/exploittiii policies in many environments. Humplik et al. (2099) proposed an alternative approach, Task Inference (TI), which trains a full RNN policy with the current task prediction an auxiliary loss. TI avoids the suboptimality of PTE/TS by not constraining the structure of the exploration/explotatiio policy. However, in TI, the task descriptors are used targets and not inputs, so TI focuses on reconstructing even irrelevant features of the task descriptor and it does not leverage the faster learning of informed policies.
In this paper, we introduce IMPORT (Inforedd POlicy RegularizaTion), a novel policy architecture for efficient online adaptation that combines the rich expressivity of RNNs with the efficient learning of informed policies. At train time, a shared policy head receives as input the current observation, together with either (learned) embedding of the current task, or the hidden state of an RNN such that the informed policy and the RNN policy are learned simultaneously. At test time, the hidden state of the RNN replaces the task embedding, and the agent acts without having access to the current task. This leads to several advantages: 1) IMPORT benefits from informed policy to speed up learning: 2) it avoids to reconstruct features of the task descriptor that are irrelevant for learning; and consequence, 3) it adapts faster to unknown environments, showing better generalization capabilities.
We evaluate IMPORT against the main approaches to online adaptation environments that require sophisticated exploration/exploitttion strategies. We confirm that TS suffers from its limited expres- sivity, and show that the policy regularization of IMPORT significantly speeds up learning compared to TI. Moreover, the learnt task embeddings of IMPORT make it robust to irrelevant or minimally informative task descriptors, and able to generalize when learning few training tasks.
2 SETTING
Let M be the space of possible tasks. Each μ M is associated an episodic μ-MDP Mμ (S,A, "r.,""",, 7) whose dynamics p" and rewards yl are task dependent, while state and action spaces are shared across tasks and 7 is the discount factor. The descriptor μ can be simple id (μ E N) or a set of parameters (n E R
When the reward function and the transition probabilities are unknown, RL agents need to devise a strategy that balances exploration to gather information about the system and exploitation to maximize the cumulative reward. Such a strategy can be defined as the solution of a partially observable MDP (POMDP), where the hidden variable is the descriptor of the MDP. Given trajectory T, = ($1, 01, F1. $t-1, a,-1, Ft-1- 1. s1). a POMDP policy π ((alTT)) maps the trajectory to actions. In particular, the optimal policy in POMDP is history- -exeendee policy that uses 7, to construct a belief state b, which describes the uncertainty about the task at hand, and then maps to the action that maximizes the expected sum of rewards (e.g. Kaelbling et al. 1998.. In this case,
2

---

Under review as a conference paper at ICLR 2021
maximizing the rewards may require taking explorative actions that improve the belief state enough so that future actions are more effective in collecting reward.
The task is sampled the beginning of episode from distribution g(H). After training, the agent returns a policy ((alT)) that aims at maximizing the average performance across tasks generated from q, i.e.,
hr
E (p)q(p)
where the expectation is taken over full-episode trajectory and task distribution g, and |r] is the length of the trajectory. The objective is then to find an architecture for и that able to express strategies that perform the best according to Eq. and, at the same time, can be efficiently learned even for moderately short training phases.
At training time, we assume the agent has unrestricted access to the task descriptor μs. Access such a task descriptor during training is a common assumption in the multi-task literature and captures large variety of concrete problems. It can be of two types: to (physical) parameters of the environment/agent (for instance, such features maybe available in robotics, or when learning on simulator) Yuu al. 2018 Mehta et al. 2019 Tobin al. 2017 7). ii) It can be single task identifier (i.e an integer) which is less restrictive assumption (Choi et al. 2001 Humplik et al.) 2019) and corresponds different concrete problems: learning set of M training levels in a video game, learning to control M different robots or learning to interact with M different users.
vector of features corresponding
RELATED WORK AND CONTRIBUTIONS
In this section, we review how the online adaptation setting has been tackled in the literature. The main approaches depicted in Fig. 2 We first compare the different methods in terms of expressiveness, and whether they leverage the efficient learning of informed policies. We then discuss learning task embeddings and how the various methods deal with unknown irrelevant task descriptors. The last subsection summarizes our contributions.
Evaluation of RL agent in Metaa-einforremeet Learning. The online adaptation evaluation setting is standard in the Meta-RL literature YYu et al.) 2017: [Humpli al. 2019) but is not the only way to evaluate agents on unseen tasks in the meta RL literature. Indeed, several works have considered that given a new task, agent is given amount of "free" interactions episodes or steps to perform system identification, then is evaluated on the cumulative reward on one (Bharadhwaj al. 2019:; Rakelly et al.) 2019) or several execution episodes (Liu et al., 2020). This is different to what we study here where the agent has to identify the task to solve and solved it within one episode, the reward of the agent being considered during all these steps.
Online Adaptation with Deep RL. In the previous section mentioned that the best strategy corresponds to the optimal policy of the associated POMDP. Since the belief state by is a sufficient statistic of the history Tr, POMDP policies takes the form x < ( T[aess Se. bc). While it is impractical to compute the exact belief state even for toy discrete problems, approximations can be learnt using Recurrent Neural Networks (RNNs) (Bakker) 200T: Heess et al. 2015). RNN-based policies are trained to maximize the cumulative reward and do not leverage task descriptors at train time. While this class of policies can represent rich exploratory strategies, their large training complexity makes them impractical.
In order to reduce the training complexity RNN policies, existing strategies have constrained the set of possible exploratory behaviors by leveraging privileged information about the task. Probe- Then-Exploit (PTE) (e.g. Zhou et al. 2019) works in two phases. First, it executes a pure exploratory policy with the objective of identifying the underlying task μ, i.e. maximizing the likelihood of the task, then runs the optimal policy associated to the estimated task. Both the probing and the informed policies are learned using task descriptors, leading much more efficient training process. PTE has two main limitations. First, similarly to explore-then. commit approaches in bandits (c.g. Garivier et al. 2016), the exploration be suboptimal because it is not reward-driven: valuable time is
3

---

Under review as a conference paper at ICLR 2021
wasted to estimate unnecessary information. Second, the switch between probing and exploiting is hard to tune and bebblem-dependent.
Thompson Sampling (TS) (Thompson 1933) leverages randomization to mix exploration and ex- ploitation. Similarly to the beliet state of an RNN policy, TS maintains a distribution over task descriptors that represents the uncertainty on the current task given 7. The policy samples a task from the posterior and executes the corresponding informed policy for several steps. Training is limited to learning informed policies together with a maximum likelihood estimator to map trajectories to distributions over tasks. This strategy proved successful in a variety of problems (e.g. Chapelle & Li, 2011 Osband & Roy, 2017). However, as shown Fig. I TS cannot represent certain probing policies because it is constrained to executing informed policies. Another drawback of TS approaches that the re-sampling frequency needs to be carefully tuned.
The Task Inference (TI) approach (Humplik et al. 2019) is RNN trained to simultaneously learn good policy and predict the task descriptor jr. Denoting by m H - the mapping from histories to latent representation of the belief state (Z C RR the policy x (a, 2x) selects the action based on the representation 2; = mm(r)) constructed by the RNN. During training, 7t is also used to predict the task descriptor μ., using the lask-identiiee monkiideetification module Z - M. The overall objective is:
|
E
+ BE
E g(=)) ]
where E(ju, g(=)) is the log-likelihood of under distribution g(=;). The auxiliary loss is meant to structure the memory of the RNN m rather than be an additional reward for the policy, so training is
done by ignoring the effect of to m. Humplik et al.] (2019) proposed two variants, AuxTask and TI, described in Fig. (b) and (c).
when computing the gradient of the auxiliary loss with respect
In TI, the gradient of the policy sub-network is not backpropagated through the RNN (the dashed
green arrow in Fig. 2c. and the policy subnetwork receives the original state features as additional input. For both AuxTast and TI, the training of in TI is purely reward-driven, so they do suffer from the suboptimality of PTE/TS. However, in contrast to PTE/TS, they do not leverage the smaller sample complexity of training informed policies, and the auxiliary loss is defined over the whole value of while only some dimensions may be relevant to solve the task.
Learning Task Embeddings While in principle the minimal requirement for the approaches above is access to lask identifiers, i.e. one-hot encodings of the task, these approaches are sensitive to the encoding on task descriptions, and prior knowledge on them. In particular, irrelevant variables have significant impact on PTE approaches since the probing policy aims at identifying the task. For instance, an agent might waste time reconstructing the full μ when only part of yl is needed to act optimally w.r.t the reward. Moreover, TS, TI and Tux are guided by a prior distribution over that has to be chosen by hand to fit the ground-trutt -ruth distribution of tasks. |Rakelly et al. (2019) proposed use factored Gaussian distribution over transitions as task embedding architecture rather than a RNN.
Several approaches have been proposed learn task embeddings Guppa et al. 2018 Rakelly et al. 2019: Zintgraf et al. 2019: Hausman et al. 2018). The usual approach is to train embeddings of task identifiers jointly with the policies. Humplik et al. (2019) mentions using TI with task embeddings, but the embeddings are pre-trained separately, which requires either additional interactions with the environment or expert traces. Nonetheless, we show in our experiments that TI can be used with task descriptors, considering task prediction as a multiclass classification problem.
Summary of the contributions As for RNN/TI, IMPORT learns an RNN policy to maximize cumulative reward, with no decoupling between probing and exploitation. As such, our approach does not suffer from scheduling difficulties instrinsic to PTE/TS approaches. On the other hand, similarly to PTE/TS and contrarily to RNN/TI, IMPORT leverages the fast training of informed policies through a joint training of an RNN and an informed policy. In addition, IMPORT does not rely on probabilistic models of task descriptors. Learning task embeddings makes the approach robust to irrelevant task descriptors contrary to TI, makes IMPORT applicable when only task identifiers are available and able to better generalize when few training tasks are available."
4

---

Under review as a conference paper at ICLR 2021
TS
AuxTask
TI
IMPORT
dt-1
CC--
MI.P
MEP
MIP
F (pes)
GRU
(T)
m((T))
UMLP
MLP
Pe (f+)
MLP
MIP
MLP
MLP
(
(xx)]
(
((t))
i (81; Rt]
d(M)
MLE wit.. (Mt, mmax roward
MLE w.r.t (JAL)) reward
MLE w.E.t (JAx) max reward
max reward
Figure 2: Representation of the different architectures. IMPORT is composed of two models sharing parameters: The (black+blue) architecture the informed policy тn optimized through (B) while the (black+red) architecture is the history-basee policy тH (used at test time) trained through (A) (A)+(C).
Algorithm 1 IMPORT Training
Initialize 0, wi, randomly for 1,.... K do
if k is odd then
Collect M transitions following тH
Update a, w and the parameters of the value function of (A) based on objective (A) (C)
else
Collect M transitions following T.
Update 0, and the parameters of the value function of (B) based on objective (B) + (C)
end if
end for
4 METHOD
In this section, we describe the main components the IMPORT model (described Fig. 2 as well the online optimization procedure and additional auxiliary loss further speed-up learning.
Our approach leverages the knowledge of the task descriptor j and informed policies to construct a latent representation of the task that is purely reward driven. Since is unknown at testing time, use this informed representation to train a predictor based on a recurrent neural network. To leverage the efficiency of informed policies even this phase, propose architecture sharing parameters between the informed policy and the final policy such that the final policy will benefit from parameters learned with privileged information. The idea to constrain the final policy stay close to the informed policy while allowing it to perform probing actions when needed to effectively reduce the uncertainty about the task. We call this approach InforMed POlicy Regulariza Tion (IMPORT).
Formally, we denote by (ag Ss, (1) and TH (ealTz) the informed policy and the history-dependerr (RNN) policy that is used test time. The informed policy HH = 6o fa is the functional composition of fu and ó. where f. M the action based on the latent representation. The idea is that fy (1) captures the relevant information contained in y while ignoring dimensions that are not relevant for learning the optimal policy. This behavior obtained by training Tr directly to maximize the task reward n*
projects μ in latent space Z S R and ő
Z - selects
While Tú leverages the knowledge of μ at training time, тн acts based on the sole history. To encourage ти to behave like the informed policy while preserving the ability to probe, лH and Fw share Φ. the mapping from latent representations to actions. We thus define as тH = x fa where fH H. Z encodes the history into the latent space. By sharing the policy head ø, the approximate belief state constructed by the RNN is mapped to the same latent space as μ. When the uncertainty about the task is small, #H then benefits from the joint training with π..
5

---

Under review as a conference paper at ICLR 2021
Figure 3: Maze 3D. The goal is either located at the blue or the red box. When the back wall (i.e. not observed in the leftmost image) has wooden texture, the correct goal is the blue box, whereas if the texture is green, the red box is the goal.
More precisely, let 0,w, a the parameters of Φ, ftr and fn. respectively, so that Th ( (ax) 11) = d f = de (ar st, f (4)) and TH (aelm) 8 o f' = - (a, st, fi (Ta)). The goal of IMPORT to maximize over 0, wi a the objective function defined in Eq.
Fr|
E
8
+
E
Tu
3E
Σ D (f.(W))
(3)
t=1
(A)
(B)
(
Speeding Up the Learning. The optimization of (B) in Eq. produces reward-driven latent representation of the task through fu- In order to encourage the history-based policy to predict a task embedding close to the one predicted by the informed policy, we augment the objective with an auxiliary loss (C) weighted by 3 > 0. D is the squared 2-norm in our experiments. Note that because we treat the objective (C) as an auxiliary loss, only the average gradient of D with respect to fH is backpropagated, ignoring the effect of fu on тH The expectation of (C) is optimized over w.0 and xz respectively used to compute (A) and (B).
trajectories generated using TH
Optimization. IMPORT is trained using Advantage Actor Critic (A2C) (Mnih et al. 2016)) with generalized advantage estimation (GAE) (Schulman et al. 2015). There are two value functions one for each objective (A) and (B). The algorithm is summarized in Alg. [ Each iteration collects a batch of M transitions using either тH or n If the batch is sampled according to TH, we update with A2C-GAE the parameters of the policy and according to both objectives (A) and (C), as well as the parameters of the value function associated to objective (A). If the batch is sampled according to Tas update with A2C-GAE the parameters of the policy and @ according to both objectives (B) and (C), as well as the parameters of the value function associated to objective (B).
5 EXPERIMENTS
We performed experiments five environments. The CartPole and Acrobot environments from OpenAl Gym (Brockman et al.) 2016), where the task descriptor j represents parameters of the physical system, e.g., the weight of the cart, the size of the pole, etc. The dimension of µ is 5 for Cartpole and 7 for Acrobot. The entries of fi are normalized in [-1, 1] and sampled uniformly. These environments provide basic comparison points where the optimal exploration/exploitttion policy is relatively straightforward. since the dynamics can be inferred from a few actions. The Bandit environment standard Bernoulli multi-armed bandit problem with K arms. The vector μ IR* denotes the probability of success of the independent Bernoulli distributions. Each dimension of μ sampled uniformly between 0 and 0.5, the best arm randomly selected and associated probability of 0.9. An episode is 100 arm pulls. At every timestep the agent is allowed to pull an arm
In our implementation, the value network shared takes as an input either fu(() or (7:). In practice, data collection is multihhreaded. We collect 20 transitions per thread with 24 to 64 threads depending the environment, based available GPU memory

---

Under review as a conference paper at ICLR 2021
N
20
N 100
= 10
N = 20
N 100
RNN
(3.8(4.8)
(1.3
87.5(0.2)
Using ja at train time
Using task identifier at train time
IMPORT
94.4(0.7)
94.8(0.8)
95.3(0.4)
92.8(0.8)
95.5(0.4)
95.1(1.0)
AuxTask
91.0(0.7)
92.0(1.9)
92.6(0.7)
90.5(1.8)
91.2(1..)
94.3
TI
91.5(0.6)
94.4(0.4)
94.6(0.3)
90.8(0.5)
90.7(1.2)
97.0(0.2) 91.0(0.5)
TS
88.7(4.2)
87.3
91.3(1.7)
85.9(2.4)
90.1(1.3)
Table I: CartPole with different number N of training tasks. Note that RNN does not at train time.
S S=1
S ==3
S= 5
S 5=1
S 3
S
Size of μ::
5
60
150
5
60
150
RNN
74.0(5.8)
77.0(1.1)
66.9(0.7)
Using
Using task identifier
IMPORT
89.2(0.1)
79.4(0.6)
72.7(0.1)
89.1(0.1)
79.7(0.1)
74.0(0.3)
AuxTask
84.7(1.1)
76.2(2.0)
66.9(1.2)
87.7(0.8)
78.0(0.4)
72.4(0.5)
TI
78.7(0.9)
76.5(0.3)
68.3(0.8)
85.3(0.8)
77.5(1.1)
70.9(0.2)
TS
85.9(0.5)
64.1
60.7(0.1)
84.1(2.5)
65.7(0.3)
60.2(1.0)
Table 2: Result over Tabular-MDP with S states and A 5 actions, trained over N 100 tasks.
in [I. KJ and observes the resulting reward. Although relatively simple, this environment assesses the ability of algorithms to learn nontrivial probing/exploitation strategies. The Tabular MDP environment is finite MDP with S states and A actions such that the transition matrix is sampled from a flat Dirichlet distribution, and the reward function is sampled from uniform distribution in [0, 1] as in Duan et 2 In that case, f is the concatenation of the transition and the reward functions, resulting in vector of size SZ SA. This environment is much more challenging as y is high nigh-diienssonal, there is nearly complete uncertainty on the task hand and each task reinforcement learning problem. Finally, the Maze 3D environment is 3D version of the toy problem depicted in Fig. I implemented using gym-miniworld (Chevalier-Botee Boisvert 2018). has three discrete actions (forward, left, right) and the objective is to reach one of the two possible goals (see Figure 15 in appendix), resulting in a reward of +1 (resp. -1) when the correct (resp. wrong) goal is reached. The episode terminates when the agent touches box or after 100 steps. The agent always starts at a random position, with a random orientation. The information about which goal to reach at each episode is encoded by the use of two different textures the wall located the opposite side of the maze w.r.t. the goals. This domain allows to evaluate the models when observations are high dimensional (3 x 60 60 RGB images). The maximum episode length is 100 on CartPole, Bandit, Rabular-MDP and Maze3D, and 500 on Acrobot. To evaluate the ability of IMPORT and the baselines to deal with different types of task descriptors μ., we also perform experiments on CartPole and Tabular-MDP in the setting where μ is only a task identifier (i.c., a one-hot vector representing the index of the training task) which is very weak supervision available at train time.
We compare to previously discussed baselines. First, a vanilla RNN policy (Heess et al. 2015) using GRUs that never uses μ.. Second, we compare to TS. TI and AuxTask, with μ only observed train time, similarly to IMPORT. For TS, at train time, the policy conditions on the true μ,, whereas at test time, the policy conditions on an estimated / resampled from the posterior every steps where A E {1, 10, 20). On bandits, UCB (Auer 2002) with tuned exploration parameters
topline.
Implementation details Contrarily to IMPORT, TS, TI and AuxTask are based on maximizing the log-likelihood of μ. When using informative task descriptors (i.e. vector of real values), the log- likeeliodooo uses a Gaussian distribution with learnt mean and diagonal covariance matrix. For the bandit setting, we have also performed experiments using beta distribution which may be more relevant for this type of problem. When using task identifiers, a multinomial distribution is used. All approaches are trained using A2C with Generalized Advantage Estimation (Mnih et al. 2016
7

---

Under review as a conference paper at ICLR 2021
16
assertte excente OX:
0.2
0.5
2.0
2.5
155
0.5
1.
(a) Bandits with K = 10, N 100
(b) CartPole with N 10
Figure 4: Test performance of IMPORT for different values of from Eq.
Schulman et al. 2015). The precise values of the aiper-parametrs. and architectures are given in Appendix [.2
All approaches use similar network architectures with the same number of hidden layers and units.
Evaluation The meta- leatnle scenario is implemented by sampling N training tasks, N validation tasks and 10, 000 test tasks with no overlap between task sets (except in Maze3l where there is only two possible tasks). Each sampled training task given a unique identifier. Each model is trained on the training tasks, and the best model is selected on the validation tasks. We report the performance on the test tasks, averaged over three trials with different random seeds, corresponding to different tasks. Training uses discount factor, but for validation and test, we
sets of train/validation/// compute the undiscounted cumulative reward on the validation/test tasks. The learning curves show test reward a function of the environment steps. They the average of the three curves associated to the best validation model of each of the three seeds used to generate different tasks sets.
Overall performances. IMPORT performs better than its competitors in almost all the settings. For instance, on CartPole with 10 tasks (see [T) our model reaches 94.4 reward while TI reaches only 91.5. Qualitatively similar results are found on Acrobot (Table 5 in Appendix), as well as on Bandit with 20 arms (Table 3). even though AuxTask performs best with only 10 arms. IMPORT particularly shines when j encodes complex information, as on Tabular-MDP (see Table 2 where outperforms all baselines in all settings. By varying the number of training tasks on CartPole and Acrobot, we also show that IMPORT's s advantage over the baselines larger with fewer training tasks. In all our experiments, expected, the vanilla RNN performs worse than the other algorithms.
Sample Efficiency. Figure 5 shows the convergence curves on CartPole with 10 and 100 training tasks and are representative of what we obtain on other environments (see Appendix). IMPORT tends to converge faster than the baselines. We also observe a positive effect of using the auxiliary loss (3 0) on sample efficiency, in particular with few training tasks. Note that using the auxiliary loss is particularly efficient in environments where the final policy tends to behave like the informed on.
Influence of μ.. The experiments with uninformative μ (i.e., task identifiers) reported in Table I and 2 for CartPole and Tabular MDP respectively show that the methods are effective even when the task deseriptors do not include any prior knowledge. In the two cases, IMPORT
these tasks descriptors to generalize well. Moreover, experimental results on CartPole (Fig. TT and Tabular MDP (Fig. 17 suggest that when μ is vector of features (and not task identifier only) it improves sample e efficiency but does not change the final performance. This can be explained by the f'ct that informed policies are faster to learn with features in μ since, in that case, μ is capturing similarities between tasks. Equivalent performance of IMPORT on both types of task descriptors is observed and shows that our method can deal with different (rich and weak) task descriptors. We further analyze the impact of the encoding of μu the models, by using non-linear projections of the informative to change the shape of the prior knowledge. Figure 55 shows the learning curves of TI and IMPORT on CartPole with task identifiers, the original μ and polynomial expansions of μ of order 2 and 3. resulting in 21 and 56 features. IMPORT's s task embedding approach is robust to the encoding of μ. while TT's likeliiiddood approach underperforms with the polynomial transformation.
Task embeddings. To have a qualitative assessment of the task embedding learnt by IMPORT, we consider bandit problem with 10 arms and embedding dimension 16. Figure 6 shows the clusters of task embeddings obtained with t-SNE (Maaten & Hinton 2008). Each cluster maps to an optimal arm, showing that IMPORT structures the embedding space based on the relevant information. In
8

---

Under review as a conference paper at ICLR 2021
100
Inpert
0.3
Tarestacs
Timmsdess
H23T
(a) CartPolc N 10
(b) CartPole N 100
64
mpaat
80
Incentia recorbine
001
escontplly
Fipey, 11
HOSt
(c) Effect transforming μ (CartPole, 20).
(d) Maze 3D.
Figure 5: Learning curves on CartPole (a and b) and Maze3D ( test tasks. Figure studies the impact of the structure of the task descriptor on the performances of TI and IMPORT in CartPole.
K 10
K 20
IMPORT
77.5(0.2)
56.6(0.1)
AuxTask (Gaussian)
78.7(0.4)
50.5(1.6)
AuxTask (Beta)
78.2(0.7)
37.1(0.6)
RNN
73.6(0.7)
32.1(1.2)
(Gaussian)
78.7(1.6)
41.4(2.4)
TI (Beta)
79.5(0.1)
53.3(2.4)
TS (Gaussian)
50.4(0.4)
38.8(2.0)
TS (Beta)
41.3(1.5)
36.3(1.1)
UCB
78.5(0.3)
68.2(0.4)
Figure 6: Task embeddings learnt on Bandit (10 arms). Colors indicate the best arm.
Table Bandits performance for K and K 20 arms, with N 100 training tasks.
addition, we have studied the influence of the 3 hyperparameter from Eq. 3 (in Fig. 4 and Section Dj shows that the auxiliary loss helps speed-up the learning process, but necessary achieve great performance.
High dimensional input space. We show the learning curves on the Maze3D environment in Figure 5d IMPORT is succeeding in 90% of cases (reward R 0.8), while TI succeeds only in 70% of cases. This shows that IMPORT is even more effective with high-dimensional observations (here, pixels). IMPORT and TI benefit from knowing fı at train time, which allows them to rapidly identify that the wall texture behind the agent is informative, while the vanilla RNN struggles and reaches random goals. TS is not reported since this environment is typical failure case as discussed in Fig
Additional results. Appendix C. we show that IMPORT outperforms by a larger margin when the task embedding dimension is small. We also show that IMPORT outperforms its competitors in dynamic environments, i.c., when the task changes during the episode.
CONCLUSION
We proposed a new policy architecture for meta reinforcement learrning. The IMPORT model is trained only the reward objective, and leverages the informed policy to discover effective trade-off's between exploration and exploitation. is thus able to learn better strategies than Thompson Sampling approaches, and faster than recurrent neural network policies and Task Inference approaches.

---

Under review as a conference paper at ICLR 2021
REFERENCES
Peter Auer. Using confidence bounds for exploitationn-eploraiion trade-offs. Journal of Machine Learning Research, 3(Nov):397-42) 4229, 2002.
Bram Bakker. Reinforcement learning with long short-term memory. In Proceedings of the 14th International Conference on Neural Information Processing Systems: Natural and Synthetic, NIPS'01, 01, pp. 1475-1482 Cambridge, MA, USA, 2001. MIT Press.
M. G. Bellemare, Y. Naddaf, J. Veness, and M. Bowling. The arcade learning environment: An evaluation platform for general agents. Journal of Artificial Intelligence Research, 47:253-279, Jun 2013. ISSN 1076-9757. doi: 10. 1..16133 1...66///airr9122 U
http:/ /tt:://xxx.o /xx doi org.1.....0..0000.......... httpp://o...dd..ooo/........1
jair. 3912
Homanga Bharadhwaj. Shoichiro Yamaguchi, and Shin ichi Maeda. neural-policy generalization and adaptation, 2019.
Manga: Method agnostic
Greg Brockman, Vicki Cheung, Ludwig Pettersson, Jonas Schneider, John Schulman, Jie Tang, and Wojciech Zaremba. Openai gym, 2016.
Olivier Chapelle and Lihong Li. An empirical evaluation of thompson sampling. In Advances in neural information processing systems, pp. 2249-2257, 2011.
Maxime Chevalier-Boisvert. gym-miniworld environment for openai gym. https: //g thub 2018. MiniWorld licensed under Apache License 2.0.
com/maxi ma
Samuel Choi, Dit- Yan Yeung, and Nevin Zhang. Hidden-mode markov decision processes for nonstationary sequential decision making. pp. 264-287, 01 2001.
Yan Duan, John Schulman, Xi Chen, Peter L.. Bartlett, llya Sutskever, and Pieter Abbeel. RI$22S: Fast reinforcement learning via slow reinforcement learning. CoRR, abs/1611.02779, 2016. URL org/ /awprg//bbb///e 1611.11111 0026661122777
Chelsea Finn, Pieter Abbeel, and Sergey Levine. Model-Agnostic Meta-Learming for Fast Adaptation of Deep Networks. arXiv e-prints, art. 1IXiv:1703.034000 Mar 2017.
Aurélien Garivier, Tor Lattimore, and Emilie Kaufmann. On sxplore-then-comee strategies. In NIPS, pp. 784-792, 2016.
Abhishek Gupta, Russell Mendonca, Yuxuan Liu, Pieter Abbeel, and Sergey Levine. Meta- reinforcement learning of structured exploration strategies. In NeurIPS, pp. 5307-5316, 2018.
Karol Hausman, Jost Tobias Springenberg, Ziyu Wang, Nicolas Heess, and Martin A. Riedmiller. Learning an embedding space for transferable robot skills. In ICLR (Poster). OpenReview.net, 2018.
Nicolas Heess, Jonathan J. Hunt, Timothy Lillicrap, and David Silver. Memory-based control with recurrent neural networks. CoRR, abs/1512.04455,, 2015. URL 1512. 04455
Matteo Hessel, Joseph Modayil, Hado van Hasselt, Tom Schaul, Georg Ostrovski, Will Dabney, Dan Horgan, Bilal Piot. Mohammad Azar, and David Silver. Rainbow: Combining improvements in deep reinforcement learning, 2017.
Jan Humplik, Alexandre Galashov, Leonard Hasenclever, Pedro A. Ortega, Yee Whye Teh, and Nico- las Heess. Meta reinforcement learning as task inference. arXiv e-prints, art. arXiv:1905.0644 May 2019.
Leslie Pack Kaelbling, Michael L. Littman, and Anthony R. Cassandra. Planning and acting in partially observable stochastic domains. Artif. Intell., 101(1-2):99-134, 1998.
Alessandro Lazaric. Transfer in reinforcement learning: a framework and a survey. In Reinforcement Learning, pp. 143-173. Springer, 2012.
10

---

Under review as a conference paper at ICLR 2021
Evan Zheran Liu, Aditi Raghunathan, Percy Liang, and Chelsea Finn. Explore then execute: Adapting without rewards via factorized meta-reinforcement learning, 2020.
Laurens van der Maaten and Geoffrey Hinton. Visualizing data using t-sne. Journal of machine learning research, 255x))2577- 2008.
Bhairav Mehta, Manfred Diaz, Florian Golemo, Christopher J. Pal, and Liam Paull. Active domain randomization, 2019.
Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Alex Graves, Ioannis Antonoglou, Daan Wierstra, and Martin Riedmiller. Playing atari with deep reinforcement learning, 2013.
Volodymyr Mnih, Adrià Puigdomènech Badia, Mehdi Mirza, Alex Graves, Timothy P. Lillicrap, Tim Harley, David Silver, and Koray Kavukcuoglu. Asynchronous Methods for Deep Reinforcement Learning. arXiv e-prints, art. arXiv:1602.01 783, Feb 2016.
Ian Oshand and Benjamin Van Roy. Why is posterior sampling better than optimism for reinforcement learning? In ICMI., volume 70 of Proceedings Machine Learning Research, pp. 2701-2710.
PMLR, 2017.
lan Osband, Charles Blundell, Alexander Pritzel, and Benjamin Van Roy. Deep exploration via bootstrapped DQN. In NIPS, pp. 4026-4034, 2016.
Ian Osband, Benjamin Van Roy, Daniel J. Russo, and Zheng Wen. Deep exploration via randomized value functions. Jourmal of Machine Learning Research, 20(124):1-62 2019. URL http:
7/jmlr. /tt.....rr//pppe//////111
html
Xue Bin Peng. Marcin Andrychowicz, Wojeiech Zaremba, and Picter Abbeel. Sim-to-Real Transfer of Robotic Control with Dynamics Randomization. arXiv e-prints, art. arXiv:1710.0653) Oct 2017.
Aniruddh Raghu, Maithra Raghu, Samy Bengio, and Oriol Vinyals. Rapid Learning or Feature Reuse? Towards Understanding the Effectiveness of MAML. arXiv e-prints, art. 1rXiv:190.009157, Sep 2019.
Kate Rakelly, Aurick Zhou, Chelsea Finn, Sergey Levine, and Deirdre Quillen. Efficient off- policy meta-reinforcement learning via probabilistic context variables. In ICML, volume 97 of Proceedings of Machine Learning Research, pp. 5331-5340. PMLR, 2019.
John Schulman, Philipp Moritz, Sergey Levine, Michael Jordan, and Pieter Abbeel. High-dimensional continuous control using generalized advantage estimation. arXiv preprint 15Xiv:1506.024335 2015.
John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms, 2017.
Matthew E. Taylor and Peter Stone. An introduction inter-task transfer for reinforcement learning Al Magazine, 32(1):15-34 2011. URL htlL:htt::/wwe/ http::://www.cc. cs utw.usteass....
wttp://cs.texx.eeeoooooo wttexasseduuusees//aaa.. lab
AdAIMagl1-Taylo.
Yee Whye Teh, Victor Bapst, Wojciech Marian Czarnecki, John Quan, James Kirkpatrick, Raia Hadsell, Nicolas Heess, and Razvan Pascanu. Distral: Robust multitask reinforcement learning. CoRR, ht//1707.04175 2017. BRLpttt:://axxiii httpt//:/aaxxi..oor////////////
04175
William R Thompson. On the likelihood that unknown probability exceeds another in view of the evidence of two samples. Biometrika, 25(3/4):285-294, 1933.
Josh Tobin, Rachel Fong, Alex Ray, Jonas Schneider, Wojciech Zaremba, and Pieter Abbeel. Domain randomization for transferring deep neural networks from simulation the real world, 2017.
Aaron Wilson, Alan Fern, Soumya Ray, and Prasad Tadepalli. Multi-task reinforcement learning: A hierarchical bayesian approach. In Proceedings of the 24th International Conference on Machine Learning, ICML '07. pp. 1015-1022, New York, NY. USA, 2007. Association for Computing Machinery. ISBN 17815995339333.
doi: 10.1145/127346.. 12736244 URL https:
10. 05.1145/127396..111 1222292211122422
11

---

Under review as a conference paper at ICLR 2021
Wenhao Yu, Jie Tan, C. Karen Liu, and Greg Turk. Preparing for the unknown: Learning universal policy with online system identification, 2017.
Wenhao Yu, C. Karen Liu, and Greg Turk. Policy transfer with strategy optimization, 2018.
Wenxuan Zhou, Lerrel Pinto, and Abhinav Gupta. Environment probing interaction policies. 2019.
Luisa Zintgraf, Maximilian Igl, Kyriacos Shiarlis, Anuj Mahajan, Katja Hofmann, and Shimon Whiteson. Variational Task Embeddings for Fast Adaptation in Deep Reinforcement Learning. arXiv e-prints, 2019.
12

---

Under review as a conference paper at ICLR 2021
A THE IMPORT ALGORITHM
The algorithm is described in details Algorithm In our implementation, the value function network used for (A) and (B) is the same, i.c. shared. We specialize the input, i.c. for (A) the input will be (st. fu (T)) and (st, f (14) for (B).
Algorithm 2 Details of IMPORT Training
Initialize 0, arbitrarily
Hyperparametter::
Number of iterations
K.
Number of transitions per update steps
M.
discount factor y. GAE parameter "GAE. Adam learning rate 9 weighting of the (C) objective weighting of the entropy objective An, weighting of the critic objective Xa for K do if is odd then
Optim = Adam(n)
Collect transitions according to мy in buffer By- else
Collect M transitions according to xy in buffer By. end if
8, 8w. dg 0,0,0
R + compute.gae.return Bu, YGAE)
R + compute gae returns TGAB)
M0 += Bul EyeBy Σt [R - Viloi Vww log
80. 13 2
Σ (T +
8
Bur >
28
BB I 23) - fol (
1Byl EieBn E (R Vo (o
80.0 + B EbED, Σt= (R - Vo n4) Vo.- log n. (o.1ot.
80. + [B
Ay
ELER, ΣT Vu, (a)
8
| B. Et [R VV)
2hg
0 + Optim (0, 60)
w + Optim (w, ow.
Optim(aa d.)
Optim (v, 6w)
end
B IMPLEMENTATION DETAILS
B. DATA COLLECTION AND OPTIMIZATION
We focus on on-policy training for which we use the actor-eritic method A2C (Mnih et al. 2016) algorithm with generalized advantage estimation. We distributed execution accelerate experience collection. Several worker processes independently collect trajectories. As workers progress, shared replay buffer filled with trajectories and an optimization step happens when the buffer's capacity bs is reached. After model updates, replay buffer is emptied and the parameters of all workers are updated guarantee synchronisation.
B.2 NETWORK ARCHITECTURES
The architecture of the different methods remains the same in all our experiments, except that the number of hidden units changes across considered environments and we consider convolutional neural networks for the Maze3d environment. A description of the architectures of each method is given Fig. 2
Unless otherwise specified, MLP blocks represent single linear layers activated with tanh function
and their output size hs. All methods aggregate the trajectory into an embedding 2, using a GRU with hidden size hs. Its input is the concatenation of representations of the last action a:-1 and
13

---

Under review as a conference paper at ICLR 2021
HPs
CartPole
Acrobot
Bandits
TMDP
Maze3d
E
16
128
16
16
32
Tr
20
20
20
20
20
h8
16
32
16
64
32
h8p
{2,4, 8, 16}
{2. 4, 8, 16}
16
{16,
{2, 16, 32}
?
0.95
0.95
0.90
0.90
0.90
An
{le
1e-2 2,
7GAE clip gradient
{0.0, 1.0)
40
{le 3 3e 4
7
A
{1., 1E 1e {le 1e 2,0
B
Table Hyperparameters tested per environments. each training epoch, we run our agent on E environments in parallel collecting T'r transitions on each of them resulting in batches of M = TTTr transitions.
current state obtained separately. Actions encoded one-hot vectors. When episodes begin, we initialize the last action with a vector zeros. For bandits environments, the current state corresponds to the previous reward. TS uses the same GRU architecture to aggregate the history into Zx-
All methods use softmar activation obtain probability distribution over actions.
The use of the hidden-state differs across methods. While RNNs only use as an input to the policy and critic, both TS and TI map 2, to a belief distribution that is problem-specific, e.g. Gaussian for control problems, Beta distribution for bandits, and multinomial distribution for Maze environments. For instance, 2₁ is mapped to a Gaussian distribution by using two MLPs whose outputs of size |n| correspond to the mean and variance. The variance values are mapped to [0, 1] using sigmoid activation.
and CartPole-tas
IMPORT maps 2, an embedding SH. whereas the task embedding fa is obtained by using a tonh-activated linear mapping of μt. Both embeddings have size hsys tuned by cross-validation onto a set of validation tasks. The input of the shared policy head is the embedding associated with the policy use, ie. either fu when using TH fa when using fu-
For the Maze3d experiment and in all methods, we pre-process the pixel input s with three con- volutional layers (with output channels 32, stride is 2 and respective kernel sizes are 5, 5 and 4) and LeakyRel.U activation. We also use a batch-norm after each convolutional layer. The output is flattened, linearly mapped to vector of size hs and tanh-activated.
C EXPERIMENTS
In this section, explain in deeper details the environments and the of hyper-parameter considered. We add learning curves of all experiments to supplement results from Table 1 and in order study sample efficiency.
Task descriptor. Note that for CartPole and Acrobot y is normalized to be in [-1, 1)P where D is the task descriptor dimension. The task distribution q is always uniform, see the description of the environments for details. For experiments with task identifiers, we associate to each sampled task an integer value corresponding to the order of generation, and encode usong one-hot vector.
Hyperparameters. Hyperparameter ranges specified in Table 4 For TS, consider sampling from the posterior dynamics distribution every steps with & E [1,5, 5, 10, 20).
14

---

Under review as a conference paper at ICLR 2021
C.I CARTPOLE.
We consider the classic CartPole control environment where the environment dynamics change within set M (pp)
described by the following physical variables: gravity, cart mass, pole mass, pole length, magnetic force. Their respective pre-normalized domains are [4.8, 14.8), [0.5, 1.5], [0.01, 0.19), [0.2, 0.8), and [-10, 10]. The value of μ are uniformly sampled. Knowing some components of might not be required to behave optimally. The discrete action space (-1,1).
Episode length is T = 100.
Final performance and sample efficiency. Table shows IMPORT's performance marginally superior to other methods in most settings. Learning curves in Figure 7 allow analyzing the sample efficiency of the different methods. Overall, IMPORT more sample efficient than other methods in the privileged information μ setting. Moreover, the use of the auxiliary loss (B >> 0) usually speed-up the learning convergence by enforcing the RNN to quickly produce a coherent embedding. We can that only sharing parameters (3 0) already helps improving over RNNs.
8
20
Timesaaa
x027
(a) CartPole with y and N 10
(b) CartPole with TID and 10
E
15 0
19225
15
20
0.
0.
2.5
1.0
(c) CartPole with μ and N 20
(d) CartPole with TID and N 20
1hrwergpe
(c) CartPole with and N = 50
(f) CartPole with TID and N 50
8 1
Tiraacope
(g) CartPole with / and N 100
(h) CartPole with TID and N = 100
Figure Evaluation on CartPole
Non-stationary environments. We consider the non-stationary version of CarPole environment where at each timestep. there is a probability = 0.05 to sample a new dynamic p. Table & shows that the performance of IMPORT, AuxTask and TI are comparable in these settings.
15

---

Under review as a conference paper at ICLR 2021
Method
= 10
N 100
AuxTask
86.4(1.0)
93.0(0.3)
IMPORT
91.7(0.5)
92.7(0.8)
RNN
65.5(4.3)
89.5(0.6)
TI
88.2(3.9)
95.5(0.8)
TS
86.7(1.6)
92.2(0.7)
Tiranlags
Figure CartPole (non-stationary).
Figure 9: Non-stationary CartPole with N = 10
100
90
80
Test reward
70
import(16)
import(2)
import(4)
import(8)
ti(16)
ti(2)
ti(4)
20
ti(8)
0.0
0.2
0.4
0.6
0.8
1.0
1.2
Timesteps
x107
Figure 10: IMPORT and TI with different task embedding representation size on CartPole with N 20
Size of built embeddings. We now study the impact of the task embedding representation size. As can be seen from Figure 10, IMPORT's performance remains stable for different representation sizes in {2, 8, 16) whereas TT's sample efficiency decreases with this dimension.
Trajectory and task embeddings. In Figure 11 episode of the final model obtained training IMPORT with two-dimensional task embeddings on CartPole with task identifiers (left) and task embedding (u) learnt by the informed policy (right). As expected, the history embedding gets close to the task embedding after just few timesteps (left). Interestingly, task embeddings (1) are able to capture relevant information from the task. For instance, they are highly correlated with the magnetic force which very strong factor to "understand" from each new environment to control the system correctly. At the opposite, gravity is less correlated since does not influence the optimal policy whatever the gravity is, the pole on the left. then you have to go right and vere-versa.
we plot both the evolution of (T;) during an
16

---

Under review as a conference paper at ICLR 2021
Component: Magnetic force
0.5
Feature value
2
0.0
1
005
-2
Component: Gravity
3
0.5
Feature value
(a) Value of (T.) among episades steps CartPole with task identifiers. The green cir- cle the value of (p). The image shows 2 that IMPORT starts with random embed- ding, and is able discover the task embed- ding with reasonable performance few steps.
0.0
-0.5
-2
0
(b) Task embeddings (a) for Cartpole with task iden- tifiers. The color the point corresponds the value one the 'real' physics component the environment (unknown to the model).
Figure 11: Visualization of task embeddings upon Cartpole
C.2 ACROBOT
N = 10
N = 20
N = 50
N = 100
AuxTask
-189.0(54.8)
-98.3(1.8)
-103.0(8.0)
93.6(1.3)
IMPORT
-87.2(0.9)
-92.5(1.3)
-88.9(1.1)
-88.9(1.6)
RNN
-483.6(1.6)
-482.7(4.0)
-480.7(3.5)
-485.0(3.7)
-89.7(1.2)
-94.6(0.7)
-87.8(0.8)
-1..3(1.)
TS
-101.4(2.0)
-102.1(6.0)
-102.4(2.0)
-102.3(0.8)
Table Acrobot
Acrobot consists of two joints and two links, where the joint between the two links is actuated. Initially, the links are hanging downwards, and the goal is to swing the end of the lower link up to given height. Environment dynamics are determined by the length of the two links, their masses, their maximum ve- locity. Their respective pre-normalized domains are [0.5, 1.5), [0.5, 1.5], [0.5, 1.5), [0.5, 1.5], [3x, 5m] and [7T,, 11m]. Unlike CartPole, the environment is stochastic because the simulator applies noise to the applied force. The action space is (-1,0, 0, 1}. We also add an extra dynamics parameter which controls whether the action order is inverted, i.c. (1,0, 1), thus |n| 7.
Episode length 500.
IMPORT outperforms all baselines in settings with small training task sets (Figure 12 and Table and perform similarly to TI on larger training task sets.
17

---

Under review as a conference paper at ICLR 2021
-130
Kolesk
-320
332
Teredeys
(a) N 10
(b) 20
-120
M
(c) N 50
(d) N 100
Figure 12: Performance on Acrobot
C.3 BANDITS
The Bandit environment is standard Bernoulli multi-armed bandit problem with K arms. The vector μ E R-K denotes the probability of success of the independent Bernoulli distributions. Each dimension of ju is sampled uniformly between 0 and 0.5, the best arm is randomly selected and associated a probability of 0.9. Although relatively simple, this environment assesses the ability of algorithms to learn nontrivial exploration/exxloittiio strategies.
Note that is not surprising that UCB outperforms the other algorithms this setting. UCB an optimal algorithm for MAB and have optimized it for achieving the best empirical performance. Moreover, IMPORT cannot leverage correlations between tasks since, due to the generation process, tasks are independent.
We visualize the task embeddings learnt by the informed policy in 13
Figure 13: t-SNE of the task embeddings the bandit problem with K 10.
18

---

Under review as a conference paper at ICLR 2021
00
Pacoe srgnnssarr
16
ratwrt)
Inpprtt
(rpacnsent)
(a) Bandits with K 10 arms
(b) Bandits with K 20 arms
Figure 14: Learning curves on the bandit problem.
C.4 MAZE3D ENVIRONMENT
Figure 15: Maze 3D. The goal is either located the blue the red box. When the wall the opposite side of the boxes (i.e. not observed in the leftmost image) has a wooden texture, the correct goal is the blue box, whereas if the texture is green, the red box is the goal.
The Maze 3D environment (Figure 15 is a continuous maze problem implemented using gym- miniworld (Chevalier-Boisverr] 2018 8). with 3 discrete actions (forward, left, right) where the objective is to reach one of the two possible goals, resulting in reward of +1 (resp. -1) when the correct (resp. wrong) goal is reached. If box touched, the episode ends. The maze's axis range from -40 to 40, the two turn actions (left, right) modify the angle by 45 degrees, and the forward action is length move. The agent starts in a random position with random orientation. The information about which goal to reach at each episode is encoded by the use of two different textures on the wall located on opposite side the boxes. this way, the agent cannot simultaneously observe both boxes and the "informative" wall.
This environment allows to evaluate the models a setting where the observation is high dimen- sional space (3x60x60 RGB image). The mapping between the RGB image and the task target in {-1,1} 1} is challenging and the informed policy should provide better auxiliary task targets than TI thanks to the "easy" training of the informed policy.
IMPORT outperforms TI on this environment (Figure 16 in both final performance and sample efficiency.
19

---

Under review as a conference paper at ICLR 2021
0.3
import
Test reward
0.4
0.0
0.0
0.5
1.0
155
2.0
2.5
3.0
3.5
4.0
limesteps
x19
Figure 16: Learning curves on the Maze 3D environment
C.5
TABULAR MDPs
Tabular MDP (Duan et al. 2016) is MDP with S discrete states and A actions such that the transition matrix is sampled from a flat Dirichlet distribution, and the reward function is sampled from a uniform distribution in [0, 1]. The task identifier is concatenation of the transition and reward functions resulting in vector of size S + SA, allowing to test the models with high-dimensional μ.
IMPORT outperforms all baselines in all settings (Figure 17 and Table 2
12
8
8
1022
TMDP with μ and
TMDP with TID and |S|
10
0.2
0.5
1.0
25 2.0
2.
0.5
255
(c) TMDP with j and SS| 3
(d) TMDP with TID and [S]
15.0
02.5
01.3
G0.5
60.00
600
5.5
X32"
(c) TMDP with 1 and S 5
TMDP with TID and |S]
5
Figure 17: Evaluation Tabular-MDP with different parameters and task descriptors (TID stands for task identifier).
20

---

Under review as a conference paper at ICLR 2021
D IMPACT OF THE 3 HYPERPARAMETER
We study the sensibility of the 3 parameter on IMPORT. Figure 18 clearly shows the benefits using the auxiliary objective. On all but the Tabular-MDP environments, the recurrent policy successfully leverages the auxiliary objective to improve both sample efficiency and final performance for Acrobot.
8 13
importin.11
recorto
Veconto
Timasaacs
Timasnas
XO2T
(a) Bandits. K 10, N 100
(b) Bandits. 20,
100
240
Rncertoaa
Rncartte 011 espariak 31
expertt
0.2
63
63
3.2
1.2
(c) CartPole with N = 10
(d) CartPole with N 100
-300
-223
epottt on rp
-500
-120
Tirasbac
(e) Acrobot with N 10
(f) Acrobot with 100
(g) TMDP with |S]
A = 5. N 100
(h) TMDP with |S] = 3, |4| = 5, N = 100
Figure 18: Test performance of IMPORT for different 3 parameters (auxiliary supervised objective). We only report performance on informative μ task descriptors.
21


Reviews:
[{"review_id": "bqsEOPprZ0q", "reviewer": "AnonReviewer4", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "The authors propose an alternative architecture to handle explore/exploit tradeoffs in RL environments where each task instance may change in such a way that the policy needs to change in order to be optimal. Rather than using an explicit task inference process, and rather than relying on an RNN to slowly learn the distribution implicitly, the task id is observed during training instances and both an embedding and an RNN are trained, such that the two are interchangeable. Thus, during testing, the task id is not needed and only the RNN state is used to condition the policy. It is a straightforward way to include privileged information during training without imposing the burden of reconstruction.\n\nThe paper is clearly written and Fig 1 is very helpful to understanding the details of the architecture. The experiments are clearly explained.\nThe main question as a reviewer is whether the paper has significance to the community. Although it is only a small architectural contribution, the method works impressively well. It is faster to learn than Task Inference and achieves higher scores than Thompson Sampling. \n\nIt would be nice to know if the method could work in combination with other methods to quickly adapt in dynamic environments, given some labels for different features of the environments. In general, using an interchangeable embedding and RNN state is a good way to avoid the challenges of conditional architectures. The paper could be stronger if the method was framed more generally and it was shown that it could be useful on a broader range of domains that require adaptation or exploration/exploitation strategies.", "extra_scores": {}, "extra_text": {}}, {"review_id": "PRd2urm85hT", "reviewer": "AnonReviewer1", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "This paper presents a method that leverage task descriptors for multi-task learning. In the proposed method, an informed policy that takes in the task descriptor and the state is trained to maximize the expected return. At the same time, a RNN policy based on the history of states and actions is trained such that the RNN layers imitates the behavior of the feature extraction layers of the informed policy. In this way, the RNN policy is trained as if the task description is available. The experimental results show that the proposed method outperforms baseline methods that leverages the tasks descriptor.\n\nThe proposed method seems novel and the experimental results show its benefits. However, there are some unclear points. Especially, “online adaptation” described in the introduction is not clear.  I would like to ask the authors to clarify the following points:\n\n- The experiment procedure is not clear to me. I understand that the informed policy and the RNN policy are trained on training tasks, but I’m not sure how the policy is adapted for test tasks. Both informed and RNN policies are further trained on test tasks? \n\n- Does Figure 4 shows the learning curve during the training on the training tasks? If so, I recommend to add the learning curve on the test tasks to show the performance of adaptation. \n\n- In page 7, I do not clearly understand this sentence: “Each model is trained on the training tasks, and the best model is selected on the validation tasks.” Were several models trained on training tasks? If so, how many models were trained? \n\n- If I understand correctly, the policy is trained to maximized the expected return across the training tasks. If so, for clarify, I recommend to describe the expectation explicitly in Eq. (3), e.g, E_{\\mu \\sim p(\\mu)} [ E_{s \\sim p(s’|s, a), a \\sim \\pi(a|s)} [ … ] ] \n\n- Caption of table 1 “Note that RNN does not \\mu at train time.” <- something is wrong?\n\n-\tHow to sample action at the initial step in the test tasks? The RNN policy seems to require the previous action a_{t-1} to generate actions, but a_{t-1} is not available in the first time step. \n", "extra_scores": {}, "extra_text": {}}, {"review_id": "N7cQQ3JAylF", "reviewer": "AnonReviewer3", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "Summary: This paper studies the exploration in meta-RL problem, where a meta-RL agent must both explore (to reduce uncertainty about the task) and then solve the task. Prior RNN approaches can theoretically learn the optimal policy, but optimization can be challenging. Instead, this paper opts to leverage additional information in the form of task-descriptors that specify the task, which make the learning process easier, because the task-descriptor provides the information normally discovered via exploration. In contrast to other task-descriptor approaches, which use Thompson Sampling for exploration (and can be arbitrarily sub-optimal in certain tasks), or auxiliary losses for predicting the task-descriptor (which may try to predict task-irrelevant information), this paper instead proposes to condition on the task-descriptor as an input to the policy during meta-training. Then an RNN policy is trained in two ways: first, conditioned on the task-descriptor, which enables learning to solve the task without requiring exploration; and second, conditioned on its own hidden state, which eventually replaces the task-descriptor at meta-test time. This approach outperforms prior approaches on a suite of tasks.\n\nStrengths:\n- Clarity. Generally, this paper presents a clear and coherent narrative. The motivation for the approach is clear (better leveraging task-descriptors compared to prior approaches to more easily learn informed policies). And the approach itself is also quite understandable.\n- Technical soundness. Furthermore, the approach appears to be technically sound. Leveraging the task-descriptor to learned informed policies, which are easier to optimize, can clearly produce some benefits. And then addressing the issue that the task-descriptor is unavailable at meta-test time by leveraging the hidden state also seems reasonable. The additional auxiliary loss to keep the recurrent state and task-descriptor embedding close also seems technically sound.\n\nWeaknesses:\n- Experiments. My main concern is that the experiments do not clearly substantiate the claims made in the main text:\n    - The introduction of the paper claims that IMPORT “adapts faster to unknown environments, showing better generalization capabilities.” From the learning curves in Figure 4, this is not clearly the case. In 4a), 4b) and 4c), other methods seem to be learning in roughly as many samples as IMPORT requires. Notably, 4d) does show some impressive improvements, but the general claim of faster learning isn’t clearly supported by the current experiments. I was also unable to determine if the train / valid / test task splits overlapped at all, to evaluate the generalization to unknown environments part.\n    - The paper also states: “We evaluate IMPORT against the main approaches to online adaptation on environments that require sophisticated exploration/exploitation strategies.” However, the environments used in the experiments, beyond the Maze3D environment don’t appear to require sophisticated exploration. I would find the experiments more compelling if IMPORT was evaluated on more complex environments requiring sophisticated exploration, like Maze3D. In a similar vein, it’s unclear to me what the takeaways from the results on the current environments should be: i.e., what does it mean about IMPORT if it performs slightly better than other approaches on CartPole or a tabular MDP? To be clear, I find evaluation on simpler environments to be useful if they clearly illustrate a point about the method, but it’s unclear to me what that point is.\n    - One of the contributions of the approach, the auxiliary loss (C) is not clearly evaluated / ablated in the experiments. It appears that all of the experiments use the same value of $\\beta$? But it’s not explicitly stated.\n    - Finally, it would be nice if the experiments substantiate the claim that IMPORT outperforms TI by avoiding “reconstructing features in the task descriptor that are irrelevant for learning.” However, it’s not clear to me that this is occurring in the experiments. Concretely, when both approaches are given task identifiers, this problem doesn’t seem to exist, since only recovering the learning-relevant aspects is sufficient for predicting the identifier with TI. Similarly, it seems like the $\\mu$’s used in the experiments contain mostly learning-relevant information.\n    - In addition to the concerns about supporting the paper’s main claims, I have two additional concerns: First, the performance improvement from IMPORT is generally fairly modest. There’s only a 3-5% gain over prior approaches in CartPole, and the TabularMDP, and the results on the bandits are mixed, although the results on Maze3D are impressive.\n    - Second, the TS baseline seems strangely implemented, as it’s based on “maximizing the log-likelihood of $\\mu$.” It’s challenging to verify the details, because the paper states that they are in the appendix, but there is no appendix. In particular, for the bandits setting, it seems like TS should be updating a beta distribution over each arm, which I would expect to lead to stronger performance. Alternatively, it would be good to use prior approaches for TS, such as PEARL [1].\n\nGenerally, I find the proposed approach to be quite promising. This work convincingly states reasons why prior approaches (e.g., TS and task inference) sub-optimally leverage task-information. Yet, I find that the experiments insufficiently support the paper's claims, so I initially lean toward rejection.\n\nAdditional minor comments that did not affect my score:\n- Related works:\n    - The point about conditioning on a belief state with an RNN policy should probably cite [2].\n    - Thompson Sampling as an exploration policy in meta-RL should probably cite [1].\n    - [3] is also relevant to the exploration problem.\n- Several key details are missing from this paper, such as the details of the environments (e.g., what is $\\mu$ in the Maze3D task), which makes evaluating the experiments challenging. As mentioned above, these are reported to be in appendix, but there is no appendix.\n- It would be nice to know the failure mode of other approaches on the Maze3D environment.\n- The end of the setting section (Section 2) somewhat conflates the problem statement with the solution: i.e., the problem is maximizing returns, while the _proposed solution_ is “to find an architecture for $\\pi$ that is able to express strategies that perform the best according to Eq. 1“\n- “they approaches are sensitive” —> these approaches...\n\n[1] Kate Rakelly, Aurick Zhou, Deirdre Quillen, Chelsea Finn, and Sergey Levine. Efficient off-policy meta-reinforcement learning via probabilistic context variables. Mar. 2019. https://arxiv.org/abs/1903.08254\n\n[2] Zintgraf, L., Shiarlis, K., Igl, M., Schulze, S., Gal, Y., Hofmann, K., and Whiteson, S. Varibad: A very good method for bayes-adaptive deep RL via meta-learning. Oct. 2019. https://arxiv.org/abs/1910.08348\n\n[3] Liu, E. Z.; Raghunathan, A.; Liang, P.; and Finn, C. Explore then Execute: Adapting without Rewards via Factorized Meta-Reinforcement Learning. June 2020. https://openreview.net/forum?id=La1QuucFt8-\n\n======= UPDATE ========\n\nI appreciate the authors' efforts during the rebuttal period, but I still retain my initial assessment of the work.\n\nOverall, I find the proposed approach promising and easy to understand, but believe that the experiments can be improved to better substantiate the claims in this work. In particular, I believe that the benchmarks can still be more carefully chosen to better evaluate IMPORT's ability to perform sophisticated exploration. I find the 3D Maze experiment to be quite nice, as it clearly highlights a shortcoming of TS exploration, but I would find the experiments more compelling if there were additional benchmarks testing such exploration. The authors commented that exploration in meta-RL is about inferring the task to solve, which I agree with, but I think such exploration can still be made more \"sophisticated\" by requiring careful sequences of actions to lead to distant states, which reveal this task information.\n\nIn addition, several issues were raised regarding the TS baseline during the discussion period. The results in the bandits setting appear to be lower than those reported in other works, and it still seems like PEARL can be adapted to be a drop-in replacement for the TS baseline. I agree with the authors' assessment that the basic form of PEARL explores the setting with multiple episodes, but PEARL could just resample from the posterior every few timesteps, which is already what happens in the TS baseline.\n\nI do think this work should be published in the future with a more careful selection of experiments.", "extra_scores": {}, "extra_text": {}}, {"review_id": "1wZ_1msAUp", "reviewer": "AnonReviewer2", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "Summary\\\nWhen the task descriptor is available as the privileged information, the authors propose a novel method to learn the policy that can benefit from privileged information. It is reward-driven learning and yet can make use of privileged information for efficient exploration. The advantage of the proposed method is verified in the experiments.\n\nComments on the paper\\\nI think the authors show an advantage of the proposed method by some experiments, but I‘d like to further request the following things to make the paper more convincing.\n\n1. Because the proposed method needs the task descriptor, it would be good to explain what kind of tasks we can apply the proposed method. The wider the applicability of the proposed method is, more valuable the proposed method would be. \n2. In the experiments, the authors compare with TS, TI and AuxTask. But I would like to see the comparison with another task embedding method, such as Pearl, K. Rakelly, et al., \"Efficient Off-Policy Meta-Reinforcement Learning via Probabilistic Context Variables,\" in ICML, 2019. I guess RNN makes the training of the exploitation policy more difficult because its latent code dynamically changes especially when the state space is large such as Maze3D environment. On the other hand, the task descriptor would not change. So RNN may sometimes makes the training difficult. We can use other embedding architecture for $f_H$ such as used in Pearl. I also would like to note that when the parameter of the dynamics is used as the task descriptor, it becomes similar to Homanga Bharadhwaj et al., “MANGA: Method Agnostic Neural-policy Generalization and Adaptation”, in ICRA 2020.\n\n----------------------------------------------------------------\nUpdate\nThank you for the comments. But there is a misunderstanding. MANGA as well as PEARL are online methods. They just need the observed data during the episode. It can encode the observation data in online manner. I think it is not evident whether IMPORT performs better than MANGA or PERAL. I agree that RNN is general, but on the other hand, I am afraid that the internal state of RNN does not converge and usually fluctuate from time to time. It may be difficult to get a persistent policy during the episode in the same environment. I would like to encourage the authors to perform more convincing experiment and make the claim of the paper consistent with the experimental findings.", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper is borderline, as evidenced by all of the reviewer's scores.

The pros are:
- important and relevant topic
-  IMPORT is a reasonable, technically sound approach
- paper is relatively clear

The cons all lie in the experimental evaluation, and whether the experiments sufficiently back the claim that IMPORT can learn sophisticated exploration strategies and validate IMPORT's merits compared to prior algorithms. In particular:
- The choice of benchmarks does not sufficiently test the ability to explore in a sophisticated manner
- Lack of comparisons to PEARL and MANGA, which can readily be applied to the online setting
- The empirical improvements are relatively modest.

Overall, the cons slightly outweigh the pros of the paper. Indeed, no reviewer was willing to champion the paper's acceptance.
Decision: Reject