Unravel Structured Heterogeneity of Tasks in Meta-Reinforcement Learning via Exploratory Clustering
Authors: Zhendong Chu
Hongning Wang
Venue: Submitted to ICLR 2023 (2023)
Keywords: Meta Reinforcement Learning
Variational Inference

Abstract: Meta-reinforcement learning (meta-RL) is developed to quickly solve new tasks by leveraging knowledge from prior tasks. The assumption that tasks are drawn IID is typically made in previous studies, which ignore possible structured heterogeneity of tasks. The non-transferable knowledge caused by structured heterogeneity hinders fast adaptation in new tasks. In this paper, we formulate the structured heterogeneity of tasks via clustering such that transferable knowledge can be inferred within different clusters and non-transferable knowledge would be excluded across clusters thereby. To facilitate so, we develop a dedicated exploratory policy to discover task clusters by reducing uncertainty in posterior inference. Within the identified clusters, the exploitation policy is able to solve related tasks by utilizing knowledge shared within the clusters. Experiments on various MuJoCo tasks showed the proposed method can unravel cluster structures effectively in both rewards and state dynamics, proving strong advantages against a set of state-of-the-art baselines.

Under review as a conference paper at ICLR 2023
UNRAVEL STRUCTURED HETEROGENEITY OF TASKS IN META-REINFORCEMENT LEARNING VIA EX- PLORATORY CLUSTERING
Anonymous authors Paper under double-blind review
ABSTRACT
Meta-reinforceement learning (meta-RL) is developed to quickly solve new tasks by leveraging knowledge from prior tasks. The assumption that tasks are drawn IID is typically made in previous studies, which ignore possible structured het- erogeneity of tasks. The non-transferable knowledge caused by structured het- erogeneity hinders fast adaptation in new tasks. In this paper, we formulate the structured heterogeneity of tasks via clustering such that transferable knowledge can be inferred within different clusters and non-transferable knowledge would be excluded across clusters thereby. To facilitate so, we develop a dedicated ex- ploratory policy to discover task clusters by reducing uncertainty in posterior in- ference. Within the identified clusters, the exploitation policy is able to solve related tasks by utilizing knowledge shared within the clusters. Experiments on various MuJoCo tasks showed the proposed method can unravel cluster structures effectively in both rewards and state dynamics, proving strong advantages against a set of state-of saate-of-hhe-ae f--hh--art art baselines.
INTRODUCTION
Conventional reinforcement learning (RL) is notorious for sample inefficiency, which often requires millions of interactions with the environment to learn performing policy for new task. Inspired by the human learning process, meta-reinforcement learning (meta-RL)) is proposed to quickly learn new tasks by leveraging knowledge shared by related tasks (Finn et al. 2017: Duan et al. [016 Wang et al. (016]_. Extensive efforts have been put into learning and uti atilizing transferable knowl- edge in meta-RL. RL. For example, Finn et al. (2017) proposed to learn set of shared meta parameters which is used to initialize the local policy when new task comes. Duan et al. 2016) and Wang et al. 2016 trained an RNN encoder to characterize prior tasks according to the interaction history.
However, little attention has been paid the situations where some knowledge is only locally trans- ferable among tasks. All the aforemmntioned methods implicitly assume tasks have substantially shared structures, and thus knowledge can be broadly shared across all tasks. However, hetero- geneity among tasks exists in practice, if not prevails, which hampers the effectiveness of existing meta-RI. algorithms. For example, the necessary skills for the Go game can hardly be applied the Gomoku game, though both of them operate on the same chessboards. We formulate this scenario as more complicated but also more realistic meta-RL setting where tasks are originated from differ- ent distributions, i.e., tasks are clustered. As a result, some knowledge is locally transferable within clusters, but cannot be shared globally. We refer to this as structured heterogeneity among RL tasks, and explicitly model cluster structures in the task distribution to capture cluster-level knowledge
Structured heterogeneity has been studied in supervised meta-learning Yao al. 2019); but it a lot more challenging to be handled in meta-RI., where the key bottleneck is how to unravel the clustering structure in a population of RL tasks. This can be further decomposed into two key re- search questions, namely population-level structure estimation and task-level inference. Different from supervised learning tasks where static task-specific data is already provided before meta learn- ing, the observations in RL tasks are collected by an agent's interactions with the environment. As
In this paper, do not assume the knowledge in different clusters is exclusive, and thus each cluster also contain overlapping global knowledge, e.g., motor skills in locomotion tasks.

---

Under review as a conference paper at ICLR 2023
result, in addition to the original explore-exploit trade-off an RL agent needs to handle for return maximization, now a meta-RL agent also has to balance the trade-off between the need of clustering structure identification and the need of return maximization when taking actions. Oftentimes, the action benefiting return maximization does not necessarily help clustering structure identification. The situation becomes more pressing when solving new tasks at test time, where the agent needs to differentiate what is transferable at the cluster-level, task-level, or not at all, for fast adaptation.
To handle the structured heterogeneity of tasks in meta-RL, we propose a cluster-based meta-RL algorithm, called MILET: Meta reInforcement Learning via Exploratory clusTering which is de- signed to explore clustering structures of tasks and enhance fast adaptation on new tasks with cluster- level transferable knowledge. Specifically, we perform dlsterr-based posterior inference (Rao et al.) 2019 Diilokthanakul et al., (016) to infer the clustering structure of a new task. To accelerate cluster inference, we learn dedicated exploration policy that is aware of cluster structures and is trained to explore the environment to reduce uncertainty in cluster inference as fast as possible. Furthermore, we design composed reward function for it to encourage a coarse-to-line exploratory strategy. New tasks can be quickly adapted with the explored cluster information by using locally transfer- able knowledge within clusters. We compare baselines on various MuJoCo environments (Todorov et al., (2012) with cluster structures in both rewards and state dynamics. We also show our method can mitigate the sparse reward issue by sample-efficient exploration on cluster structures, which provides hints to solve a specific task. We further test our method on environments without explicit clustering structures, the results show our method can automatically discover locally transferable knowledge and benefit the adaptation.
method with rich set of state-of-the-an meta-RL
2 RELATED WORK
Task modeling in meta-learning. Task modeling is important to realize fast adaptation in new tasks in meta learning. Finn et al. (2017) first proposed the model-agnostic meta learning (MAML) aiming to learn a shared model initialization, i.c., the meta model, given a population of tasks. MAML does not explicitly model tasks, but it expects the meta model to be only a few steps of gradient update away from all tasks. Later, an array of methods extend MAML by explicitly modeling tasks using given training data in supervised meta-learning. Lee & Choi (2018) learned task-specific subspace of each layer's 's activation, on which gradient-based adaptation is performed. Vuorio et al.) (2019) explicitly learned task embeddings given data points from each task, and then used to generate task-specific meta model. Yao et al. (2019) adopted hierarchical task clustering structure, which enables cluster-specific meta model. Such design encourages the solution to capture locally transferable knowledge inside each cluster, similar to our MILET model. However, task information not explicit in meta-RL: since the true reward/state transition functions are not accessible by the agent, the agent needs to interact with the environment to collect observations about the tasks, while maximizing the return from the interactions. MILET performs posterior inference of a task's 's cluster assignment based its ongoing trajectory; better yet, it is designed to behave exploratorily to quickly identify tasks' clustering structures.
Exploration in meta-reinforcemmnt learning. Exploration plays an important role in meta-RL, as the agent can only learn from its interactions with the environment. In gradient- based meta-RL. (Finn et al. 2017), the local policy is trained on the trajectories collected by the meta policy, thus the exploration for task structure is not explicitly handled. Stadie et al.] (2018) and Rothfuss et al. (011)) computed gradients w.r.t. the sampling distribution of the meta policy, in addition to the collected trajectories. /uurumurthy et al.] (2020) learned a separate exploration policy for MAML to collect informative pre-update trajectories. Gupta et al. (2018) also extended MAML by using learnable latent variables to control different exploration behaviors. The context- meta-RL algorithms (Duan al., 2016;) Wang et al. 2016) automatically learn to trade-off exploration and exploitation by learning policy conditioned on the current context. ZZintgraf et al. (2020) explicitly provided the task uncertainty to the policy to facilitate exploration via variational inference. ZZhou et al. (2019) introduced intrinsic rewards encouraging exploration that can improve the prediction of dynamics. Zhang et al. (02T)) and Liu et al. (2021) developed a separate exploration policy by maximizing the mutual information between task ids and inferred task embeddings. However, all the aforementioned exploration methods are oblivious to the structured heterogeneity of tasks, leading to inefficient exploration and ignorance of locally transferable knowledge within clusters.
2

---

Under review as a conference paper at ICLR 2023
MILET learns the exploration policy to explicitly reduce the uncertainty in cluster inference, which is important to utilize locally transferable knowledge.
BACKGROUND
Meta-reinforcementt learning. We consider family of Markov decision processes (MDPs) p(M), where an MDP Mj ~ -(MM is defined by a tuple M; (S,A, A, Rú, T;, Two: y, H) with S denoting its state space, A as its action space, Ra(re-1 (tz 81+1) its reward function, T1(++1 8t, at) as its state transition function. To (5o) ) as its initial state distribution, Y as a discount factor, and as the horizon of an episode. The index represents the task id. which provided to agents in some works (Zhang et al. 2021] [iu et al.) 2021 Rakelly et al. (219)). Following Zintgraf al. (2020), we consider more general setting where the task id is unavailable to the agent, since becomes meaningless when we keep encountering new tasks. Tasks sampled from p(M) typically differ in the reward and/or transition functions. For an MDP, we run a trial consisting of N + episodes (Duan et al. 2016). Following the evaluation settings in previous works (Finn et al. 2017] Liu et al. 2021 Zhang et al. 2021 Rothfuss et al., 2019). the first episode in a trial is reserved as an exploration episode gather task-related information, and an agent is evaluated by the returns in the following N exploitation episodes.
Inside a trial, {sn, an Tay $n41), and T. = {sp, To...., St} denotes the whole interaction history collected before time I. In the exploration episode, an agent should form the most informative history 7 by rolling out an exploration policy To parameterized by t. In exploitation episodes, the agent executes the exploitation policy T parameterized by (in some prior work, Ty and T could be the same) conditioned on 7 and, optionally, the history collected in exploitation episodes 7 The agent's goal is to maximize the returns in exploitation episodes, which can be formally expressed as,
denote the agent's interaction with the MDP at time step Tn
NxH
I (TT T&) = EM, MM)
Σ Ra (=(T
(1)
where R2 (T + is the return of ف conditioned on 7 and T time step in task Mj. Mixture of RL tasks. In this paper, we consider more genearl and realistic setting, where different task distributions form a mixture,
P(M) Σ W. Pe(M),
(2)
where C' is the number of mixing components (i.e., clusters) and we is the corresponding weight of component such that > We = 1. Every task is sampled as follows,
Sample cluster e according to the multinomial distribution of Mul -..; wc);
Sample reward function R or transition function T or both from (MM)
The knowledge shared in different clusters could be different. For example, two clusters of target positions can exist in a navigational environment and each task is associated with particular target position. In the first cluster, the target positions concentrate in the upper part of the map, while the target positions in the second cluster fall in the bottom part. The knowledge that how the agent moves the upper part in the first cluster should not be transferred to the second cluster; but it cru- cial for different tasks in the first cluster. Such locally transferable knowledge cannot be effectively captured by previous meta-RL solutions, since they are only designed to realize globally transfer- able knowledge. Our method is able to discover and utilize such local knowledge by modeling the underline structures of tasks distributions.
METHODOLOGY
In this section, present MILET in detail. Firstly, we introduce how to estimate population- level task structures using the collected trajectories via cluster-based variational inference. Then,
Thh terms of environment, task and MDP used interchangeably in this paper.
3

---

Under review as a conference paper at ICLR 2023
Encoder
Decoder
Ag
8:
T-GRU
(x(2),,Aa) Aa)
4,
P
R
P6+1
Sample Ar.
h ai(c/ha))
DAL
C-GRU
T
8:
Pw
84+1
+ Tato
Eq.(12)
M, ~ p(M)
Exploration policy
h
0
Exploitation policy
Trans. Ti
π
h
Environments
Figure I: MILET architecture. The encoder processes ongoing trajectories and performs cluster- based task inference qo (=|c, he). The exploration policy Te is trained to explore the environment find the most certain cluster assignment c. The explored information is passed to the exploitation policy To to facilitate fast adaptation in the inferred task M
we explain the exploration policy trained by the exploration- driven reward, which is designed to quickly identify the cluster of a new task. In particular, in each task we first execute the exploratory policy to collect the coarse- -raiieed cluster information; then we adapt the task policy with the help of the inferred posterior cluster distribution obtain fine- -na-griiedd task information. The architecture of MILET is shown in Figure D
4.1 CLUSTER-BASED VARIATIONAL INFERENCE WITH CONSISTENCY REGULARIZATION
Since the reward and transition functions are unknown to the agent, we learn stochastic latent variable c; to infer the cluster assignment of current task M; PM(M). Based on C. we infer another stochastic latent variable =, carrying task-level information, i.e., reward/transition functions first drop the script in this section task as an example to illustrate our model design. We believe all information about the current task is encoded in (z, c), e.g., once know the cluster and task assignments, the dynamics of the task, i.e, the trajectories. When a new task arrives, we decode its characteristics by the posterior distribution c(Te) with respect to the interaction history up to time t. The inferred task information ze (z conditioned on c) is then input into the policy Ta/al
that define the task. For simplicity,
will only one
can perfectly reconstruct
(0//((r/53 (ar|se, ze).
However, the exact posterior p(z, c(T:)) defined by Eq, (2) is intractable. Instead, we learn an approx- imated variational posterior (a clT.2) qa (2) T.a. c)qo (c)T.)). in which we learn two dependent inference networks and collectively denote their parameters as O. On top of the inference networks, learn decoder Pa reconstruct the collected trajectories. The whole framework is trained by maximizing the following quantity.
E Bp(,, [logp(r((m)
(3)
where p is the distribution of trajectories induced by the policies { T&. and the environment,
and ++ v ) denotes all the trajectories collected in a trial, the length of which is denoted as H = (N H.... Note that we use trajectories from both exploration and exploitation episodes since the rewards and state transitions in the same task are generated by the same underlying MDP. и to simplify our notations in later discussions. Instead of optimizing the intractable objective (3). we optimize evidence lower bound (ELBO) w.r.t. the approx- imated posterior ge (2 elT.e) estimated via Monte Carlo sampling (Rao et al. 2019) (full derivation can be found in Appendix A).
We omit the dependencies
cluster-specific reconstruction loss
cluster-spocific reglaaazzoonn
ELBO, EEp
E (x. In ( [c)] (g [L(g((2)c (=(c))
(4)
categorical regularizaiooo
(C(T)((TT:)) p(c))
4

---

Under review as a conference paper at ICLR 2023
where (z(ccc) (Nuw(()) o is learnable cluster-specific prior, which is different from the simple Gaussian prior used in single-componett VAE (Kingma & Welling 2013). This prior allows the model to capture the unique structure of each cluster. Its parameters are included in w since the cluster structure is also a part of the environment. ãn is the latent variable sampled from (o(2),, T:e) NV (Ho(e, T4), Ta), T:))), using the reparameterizatioo trick (Kingma & Welling
2013]. (o(C)T.)) outputs the approximated posterior cluster distribution given T:t p(c) fixed non non-informative multinomial distribution representing the prior cluster distribution of tasks.
Similar to Zintgraf et al. 2020), the first term In (T in Eq. can be further factorized as
In (r xt = pn
In pa (8141 2c) + In pw p ai 8141, 2c)
(5)
H+-1
const. +
Σ (ra-f((s f(ss, az 8141, kc)) As ||s441 - $(81, as Ec)1
where (()) is the initial state distribution in task, and we consider it constant by assuming identical distribution of the initial states across clusters. The second and third terms are likelihood derived from decoders for transition and reward functions. A, control the effect of the state transition in variational inference. The density function of pu (s 0g Σz) and Pa (ra+. [8;, ag 81+1, &c) is difficult estimate in continuous state and action spaces. Denote the corresponding decoder outputs as f and 8, we use L2 distance to approximate the log-likelihood functions (Zhang el al. 2021).
In the inference networks go(=/Ta, c) and (c(77,, we follow Duan et al. (2016); Zintgraf et al. (2220) to encode the history T: by Gated Recurrent Units (GRUs) (Chung et al.) 2014). We propose stacked GRU structure (shown in Figure I differentiate the information for cluster and task inference in the hidden space. Specifically, we set a task GRU (T-GRU) and a cluster GRU (C-GRU), both of which encode the history Tæs. but with different levels of granularity. T-GRU set to capture fine-grained task-specific patterns in the history as it is optimized to reconstruct trajectories of a specific task. C-GRU captures more coarse-grained patterns as it is set to help T-GRU reconstruct all trajectories within a cluster. To realize this difference, the output h2 T-GRU is only provided to ge (=)ha (T,,, ha,, c), while the output ha of C-GRU is passed to both cluster inference go (c/h, (T:,)) and task inference go [=]ha(T.,, ha),e) c)- This also reflects dependency assumption about the task structure. We denote h = (ha, h }. and h is passed between episodes in trial.
Different from the static training data provided beforehand in supervised meta-lcarning settings, the trajectory data is incrementally collected by the agent in meta-RL, which brings both challenges and opportunities in obtaining the most informative information for structure identification. Firstly, inside trial, the structure inference improves more interactions are collected, which means the agent's 's belief about the task structure could change thereby. This introduces a contradiction since the inference result should stay consistent within a given task, no matter how trajectory changes over episodes. We attribute this property as in-trial consistency. It can be measured by (L((TC)TTa,) ) g(clT,,)) )), where & and to refer two arbitrary timesteps a trial and the KL divergence measures the inconsistency between two cluster posterior distributions. We realize the notion of cluster inference consistency via the following regularizer,
H Σ
4= = 11+
KL (ge
|| go(c(Ta41))
(6)
This regularizer minimizes the inconsistency between any two consecutive time steps inside a trial. The consistency of task inference implicitly enforced once the cluster inference becomes consis- tent. Furthermore, go (c|T:) is encouraged to find structures beyond the task-level by this regularizer, ie., cluster structures, as in-trial cluster inference is enforced to be the same.
Similarly, since the cluster-specific prior pa (=lc) is learnable, the task inference can become in- consistent if the prior changes drastically across training epochs. More seriously, oscillation in the inference of latent variable cause the collapse of policy training, as tasks across clusters might
In practice, we the Gumbel-sofftax trick to simplify the calculation.
5

---

Under review as a conference paper at ICLR 2023
be assigned with the same latent variable We conclude it as the prior consistency property and realize it via the following regularization,
f-= C KL P)g(2)c)), C
where (2/c(2 is target network and its parameters are the same (=(c) but updated in much slower pace. We finally conclude the objective in our cluster-based variational inference (CBVI),
H+ Σ
C(0, w) = (p(M)
ELBO, X ApCo
(8)
t=0
where dr and Ap are hyper-parameters to control the strength of two regularizers.
4.2 EXPLORATION FOR REDUCING TASK INFERENCE UNCERTAINTY
The adaptation procedure in new task can be decomposed into two sub sub-goals: (1) explore clustering structure; (2) solving the task with the inferred task information. These two sub- sub-goals are not in parallel, but the correct clustering structure identification is the prerequisite for maximizing task return. Thus, we are motivated to learn two separate policies as shown Figure One takes exploratory behaviors to collect cluster and task information, i.c., the exploration policy T The other is optimized to solve the task with the collected information, i.c., the exploitation policy Ts.
Clustering structures provide hints to solve specific tasks. We train separate exploration policy to provide good basis for task-solving by exploring the environment to quickly identify the cluster as- signment of task. We evaluate the quality of exploration using two principles. Firstly, whether the trajectory of an exploration episode can reduce the uncertainty of the cluster inference. Secondly, whether the inference result consistent. The exploration trajectory should lead consistent clus- ter, other than switching among clusters. We conclude them as certain and consistent exploration.
We design two intrinsic rewards to encourage certain and consistent behaviors. First, we the entropy of cluster inference network qe (c())))))) to measure the uncertainty of the inferred cluster. For new task, we look for trajectories that provide the most certain cluster inference. We formalize the objective as follows, omitting the subscript 0 and superseript &i for simplicity,
Σ In (ellT++1) 0
#-1
H(g(CT.N)) = -E [I = -E [[Ing((TT)) +
(9)
((cTT.))
We then define an intrinsic reward of each action by telescoping the second term similar to Zhang al. (2021): Liu et al. (202T).
(ax) = E In g(clT411 [$41;; = H(((clTa)) H (g(clTaa))
(10)
((c77a)
This reward favorites actions which can reduce the entropy of cluster inference; and therefore, trajectory leading to a consistent cluster inference preferred. To more explicitly measure the divergence between the inferred cluster distributions from two steps, we define another reward en- couraging consistent cluster inference,
re(a:) (L(G(C)T.) |l g(clTa+1))
(11)
The above two rewards are designed collect certain and consistent cluster-level information. After we locate the cluster, the exploration policy should focus on finding task-level information, which becomes easier under the help of cluster information. Intuitively, cluster is easier find, as infor- mation across individual tasks can be leveraged. This induces a coarse-to-fine exploration strategy. We define the following composed reward to encourage the coarse-to-fine exploration behaviors,
(a)))) r(a.) + 7h (a:) + (a))).((a))
(12)
where the reward provided by the environment. 7h and To are two temporal decaying functions,
7hk (1) by aq exp(=sp 8h S%
t)), (t) -be a. H- -()),
(13)

---

Under review as a conference paper at ICLR 2023
MiLEt
RL2
VariBAD
MetaCURE
PEARL
final performance
Ant Goal itpisodic 11
Ant Goal (Eploode 2)
Humansid Der lEpisede 11
Humanois Dir (Eplsode 2)
Avarage Returns
Frames
Frames
Frames
Frames
(a) Environments with clustered reward functions.
Params (Episode
Hopper
Walker Rand- (Episode
Walker Rand Parans (Episode
Average Returna
140
10
Frames
Frames
Frames
Frames
(b) Environments with clustered state transition functions.
Figure 2: Average test performance for episodes on MuJoCo environments.
where (a, 8),, are hyper- paramm-asseeters controlling the rate of decay. 7 should gradually decrease to 0, which means we encourage the policy to find a certain cluster the early stage, but it should begin to collect task information later. "e gradually increases from negative value to positive. At the early stage, a negative % encourages the policy to try different clusters. Later, a positive % enforces the policy to stick to the current cluster and focuses more on discovering task information by maximizing raw rewards.
Finally, the exploitation policy a inherits the hidden state h which encodes knowledge collected by the exploration policy, and is then trained to maximize the expected reward defined in Eq.(T). The detailed pseudo-codes of meta-train and meta-test phases for MILET are shown in Appendix
EXPERIMENTS
To fully demonstrate the effectiveness of MILET in handling structured heterogeneity in meta-RL, RL. In this section, we conduct extensive experiments to study the following research questions: (1) Can MILET achieve better performance than state-of-the-an meta-RL algorithms by handling structured heterogeneity in the task distribution?
(2) Can MILET effectively unravel cluster structures in both rewards and state dynamics? (3) Can the sparse reward issue be mitigated by exploratory clustering of tasks? (4) How does the number of clusters affect the final performance of MILET? We defer the ablation study to Appendix
5.1 REWARD AND DYNAMICS CLUSTER ENVIRONMENTS
Environment setup. We evaluate MILET on two continuous control tasks with clustered reward functions, simulated by MuloCo (Todorov al. 2012). In Ant-Goal, the ant robot needs to move a predetermined goal position. We create clusters by distributing the goal positions in 4 different centered areas. In Humanoid-Dir, the human-like robot controlled move towards different tar- get directions. We create 4 clusters by distributing goal directions along 4 orthogonal directions in 2D space. We also created environments with clustered transition functions by adopting two move- ment environments Hopper- Rand -aad--annaaaaa Params and Walker-Rand -kand-Params. also simulated by MuJoCo. The physical parameters of the robot, including body mass, damping on degrees of freedom, body in- ertia and geometry friction, can be manipulated to realize different transition functions of the robot's movement. The hopper and walker robots are required to move smoothly under different parameter settings. We created 4 clusters by manipulating one of the parameters to the default parameters. The detailed procedure of task generation can be found in Appendix C
time and keep the others
Baseline setup. We compared MILET with several representative meta-RL baselines, includ- ing RL2 (Duan et al. 2016), PEARL (Rakelly et al. 2019), VariBAD (Zinggraf et al. 2020) and MetaCURE (Zhang al.) 2021). For each environment, we created 500 tasks for meta-train and hold out 32 new tasks for meta-test. We report the performance on test tasks over frames in the meta-train phase. In the meta-test phase, we executed 2 episodes in each task. For algorithms
7

---

Under review as a conference paper at ICLR 2023
NMI
Episode
Episode
Episode
Episodo 2
Frames
(a) MiLEt traces.
(b) VariBAD traces.
(c) NMI score in exploration.
Figure 3: Qualitative analysis of MILET. (a) Traces of MILET on the meta-test tasks of Ant- Ant-Goal.
Cross marks represent goal positions, and the colors represent the cluster assignments produced by MILET. The dashed lines suggest the optimal traces to the centers of ground-truth clusters. (b) Traces of VariBAD
the same meta-test tasks of Ant-Goal. The traces are in the same color as VariBAD is unaware of clusters. (c) NMI of MILET's inferred clusters in the exploration episode of meta-test tasks in four environments.
with the explicit exploration policy, i.e., MILET and MetaCURE, we run their exploration policy in the first episode and exploitation policy in the second episode. We used public algorithm implemen- tations provided by the original papers. We trained MILET via Proximal Policy Optimization (PPO) (Schulman et al. 2017) and the default cluster number to 4. PEARL and MetaCURE are based on off- off-policy algorithms (Haarnoja et al. 2018); they need less frames of data to converge in the meta-train phase. We terminated them once the algorithm was converged and reported the final performance. The reported performance is averaged over 3 random seeds. More hyper-parameee settings and implementation details can be found in Appendix D
Results and analysis. Figure 2 shows the test performance of all evaluated meta-RL algorithms. We also provide qualitative analysis in Figure 3 including visualization of the models' behaviors and the clustering performance of MILET in the exploration episode measured by the normalized mutual information score (NMI).
MILET showed significant improvement against baselines in the second episode, which is expected as it benefits from the information explored in the first episode. Interestingly, we can observe even though the first episode of MILET was reserved for exploration, it still performed comparable to other methods in all four different environment setups. In the first episode, MILET behaves ex- ploratorily find the most probable cluster of the current task, and thus its traces in Figure 77a look like spirals from the starting point. VariBAD is also designed to explore by uncertainty in task inference, but its traces are close to random walk the early stage, which is less effective. In Figure 3c we can observe the NMI scores of the MILET's inferred tasks have almost converged in 20 time steps, which means the stable cluster inference can already provide the agent necessary cluster information to gain comparable performance in the first episode. In the second episode, with the better task informa- tion, MILET is able to move towards the targets directly, showing significant improvements against baselines. MetaCURE guides the exploration by task IDs, which in fact provides more information of environment than what MILET can access. However, the exploration empowered by task IDs fo- cuses more on leveraging task-level information thus is less effective exploring coarser but useful information, e.g., clusters, especially when the task IDs independent the task structure.
task information. This also explains how MILET obtained
5.2 SPARSE REWARD ENVIRONMENTS
Environment setup. We evaluated MILET on a more challenging setting where the reward is sparse. We modified Ant-Goal and Humanoid-Dir environments such that the agent only gets pos- itive rewards within small region around target positions or target directions, otherwise 0. We denote them as Ant-Gool-PPatial and Humanoid-Dir-Sparse In Ant-Goal-Pane evaluated methods failed when rewards were too sparse. Hence, we set the reward regions to cover the initial position of the robot. With sparse reward, it becomes more important for the agent to leverage knowledge across related tasks, as the feedback within single task is insufficient for pol- icy learning. More environment details are deferred to Appendix C
we found all
Results and analysis. We present results in Figure 4 By exploring and leveraging task structures in the exploration episode, MILET outperformed all baselines with a large margin. MetaCURE is designed for exploration when reward is sparse by utilizing task ids, which indeed helped it outper- form other baselines in Mott-oal-Paatinne But such exploration is on the task-level thus unable to effectively explore cluster information to enhance task modeling. On the contrary, MILET lever-
8

---

Under review as a conference paper at ICLR 2023
MILEt
RL2
VariBAD
MetaCURE
PEARL
final performance
Ant- Partiel (Spisode 1)
Art--Goal-Partiel (Eplacda 21
Itumanoid-Or- Sparse (Eplsode 3)
Himansid-Din-pprra (Eplsode 21
Average Raturns
20
Frames
Frames
Prames
Frames
Figure 4: Average test performance for episodes on sparse reward environments.
aged relatedness among tasks with in the same cluster to bootstrap policy adaptation; and as shown in Figure 3e| cluster inference can be efficiently solved, which provides an edge for accurate task inference. These two factors contribute to MILET's better performance in the exploitation episode.
5.3 INFLUENCE OF THE NUMBER OF CLUSTERS
Performance comparison on Ant- Goal and its variant.
Table 1:
In this experiment, we study how the number of clus- ters C' influence the final performance, especially when there is a mismatch between the ground-truth number of clusters and C' set by the agent. We set C' to different values and denote it in suffixes. We
Ant-Goal
Ant-U
VariBAD
-16..619.6
-162.419.2
MILET-2
-13..317.6
-128.618.8
additionally create set of tasks on the Ant-Goal en-
MILET-4
-125.4+5.1
-113.7+4.8
vironment, where the goal positions are uniformly sampled on circle. We denote it Ant-U. In this setting, there is no explicitly created clusters of tasks.
MILET-6
-123.6+4.4
-99.745.2
MILET-8
-124.2+4.7
-117.9
MILET-I0
-128.6+5.2
-142,7=10.4
The average final returns are shown in Table T Inter- estingly, we observe MILET can perform well even though there are no explicit cluster structures in Ant-U. By looking into the detailed trajectories, ferent parts as shown in Figure 5 such that knowledge from nearby tasks can be effectively shared. VariBAD mistakenly assumed all tasks can share knowledge and thus failed seriously. When C is set smaller than the ground- truth number of clusters, MILET-2 discovers more general structures. However, transferable knowledge within such structures less, causing the performance drop. How- ever, it does not mean more clusters than necessary is helpful, as less knowledge could be shared in each cluster. Ant-U, MIL.ET-8 and-10 generate unnecessary clusters, and cluster assignments are mixed at the boundary of adjacent clusters (see visualization in Appendix G). Such inaccurate clus- ter modeling causes ineffective exploration and sharing of wrong knowledge, leading to degenerated performance, Hence, correctly inferring the number of clusters is important future work.
found MILET segmented the circle into dif-
6 CONCLUSION & FUTURE WORK
In this paper, we present MILET, a clustering-based method to unravel structured heterogeneity of tasks in meta-RL. MILET is able to discover clustered task structures in a population of RL-tasks and adapt cluster- level transferable information to new tasks. To quickly identify the clus- ter assignment of new tasks, MILET learns separate exploration policy which aims to reduce uncertainty in cluster inference. We further de- sign dedicated reward function to control the exploration between the cluster- and task-level information. MILET outperformed representative meta-RL methods by utilizing structured heterogeneity in the task dis- tribution. We further demonstrated its exploration on cluster structures greatly improves its sample efficiency in sparse reward environments.
Figure 5: Traces of MILEr-6 in exploration episodes on Ant-U. Col- ors represent the cluster assignments produced by MILET-6.
Overall, MILET sheds lights on discovering structures in the task distri- bution to boost meta-RI. algorithms. Several extensions of MILET are worth to explore in the future work. We currently assume uniform cluster prior. More complicated priors can be considered to enable new features. For example, the Gaussian process prior can be used to automatically identify number of clusters in the task distribution, and possibly detect out- of-distribution tasks in the meta-test phase. Also, MILET can be combined with skill-based RL methods (Pertsch al. 2021 Nam et al. 2022) to learn cluster-level skills, which could form new basis for meta-RL. c.g., each task can be modeled as a mixture skills and different clusters of task associate with different skill distributions.

---

Under review as a conference paper at ICLR 2023
REFERENCES
Junyoung Chung, Caglar Gulcehre, Kyunghyun Cho, and Yoshua Bengio. Empirical evaluation of gated recurrent neural networks on sequence modeling. In NIPS 2014 Workshop on Deep Learning, December 2014, 2014.
Nat Dilokthanakul, Pedro AM Mediano, Marta Garnel,, Matthew CH Lee, Hugh Salimbeni, Kai Arulkumaran, and Murray Shanahan. Deep unsupervised clustering with gaussian mixture varia- tional autoencoders. arXiv preprint arXiv:1611.02645 2016.
Yan Duan, John Schulman, Xi Chen, Peter L Bartlett, llya Sutskever, and Pieter Abbeel. 2: Fast reinforcement learning via slow reinforcement learning. arXiv preprint 1oXiv:6611.02279 2016.
Chelsea Finn, Pieter Abbeel, and Sergey Levine. Model-agnostic meta-leaming for fast adaptation of deep networks. In International conference on machine learning, pp. 1126-1133 PMLR, 2017.
Abhishek Gupta, Russell Mendonea, YuXuan Liu, Pieter Abbeel, and Sergey Levine. Meta- reinforcement learning of structured exploration strategies. Advances in neural information pro- cessing systems, 31, 2018.
Swaminathan Gurumurthy, Sumit Kumar, and Katia Sycara. Mame: Model-agnostic meta- exploration. Conference on Robot Learning, pp. 910-922. PMLR, 2020.
Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with stochastic actor. In International confer- ence on machine learming, pp. 1861-1870. PMLR, 2018.
Eric Jang, Shixiang Gu, and Ben Poole. Categorical reparameterization with gumbel-softmax. In International Conference Learning Representations, 2017. URL httphtttt:://peereereee net//form? 2ieimmid==reee dd=rkE3y8ees
Diederik Kingma and Jimmy Ba. Adam: A method for stochastic optimization, arXiv preprint 1xXiv:1412.6988 2014.
Diederik Kingma and Max Welling. Auto-encoding variational bayes. arXiv:1312.61 2013.
arXiv preprint
Yoonho Lee and Seungjin Choi. Gradient-based meta-leamiing with learned layerwise metric and subspace. In International Conference Machine Learning, pp. 2927-2936. PMLR, 2018.
Evan Z Liu, Aditi Raghunathan, Percy Liang, and Chelsea Finn. Decoupling exploration and ex- ploitation for meta-reinfforcement learning without sacrifices. In International conference ma- chine learning, pp. 6925-69299999 6935. PMLR, 2021.
Taewook Nam, Shao-Hua Sun, Karl Pertsch, Sung Ju Hwang, and Joseph J Lim. Skill-based meta- reinforcement learning. In International Conference on Learning Representations, 2022. URL
hhttss://ppenevviee net/ woo.moiiddeeeeeoo
htt.pp/neeereeeeeeeeeoorrrrr.eed JeW- Fh9bV
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, et al. Pytorch: An imperative style, high- performance deep learning library. Advances in neural information processing systems, 32, 2019.
Karl Pertsch, Youngwoon Lee, and Joseph Lim. Accelerating reinforcement learning with learned skill priors. In Conference on Robot Learning, pp. 188-204. PMLR, 2021.
Kate Rakelly, Aurick Zhou, Chelsea Finn, Sergey Levine, and Deirdre Quillen. Efficient off-policy meta-reinforcemeent leamnnng via probabilistic context variables. In International conference machine learning, pp. 5331-5340. PMLR, 2019.
Dushyant Rao, Francesco Visin, Andrei Rusu, Razvan Pascanu, Yee Whye Teh, and Raia Hadsell. Continual unsupervised representation learning. Advances Neural Information Processing Sys- fems, 32, 2019.
10

---

Under review as a conference paper at ICLR 2023
Jonas Rothfuss, Dennis Lee, Ignasi Clavera, Tamim Asfour, and Pieter Abbeel. ProMP: Proxi- mal meta-policy search. In International Conference on Learning Representations, 2019. URL
https::/cppenevieet net htt.:/awwweee....oooooooooooaaa. wtw.:oommmmm----0111100-----
John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. arXiv preprini arXiv: 2017.
Bradly C. Stadie, Ge Yang. Rein Houthooft, Xi Chen, Yan Duan, Yuhuai Wu, P. Abbeel, and Ilya Sutskever. Some considerations on learning to explore via meta-reinforcement learning. ArXiv, abs/1803.01118. 2018.
Emanuel Todorov, Tom Erez, and Yuval Tassa. Mujoco: A physies engine for model-based control. In 2012 IEEE/RSJ international conference on intelligent robots and systems, pp. 5026-5033. IEEE, 2012.
Risto Vuorio, Shao-Hua Sun, Hexiang Hu, and Joseph J Lim. Multimodal model-agnostic meta- learning via task-aware modulation. Advances in Neural Information Processing Systems, 32, 2019.
Jane Wang, Zeb Kurth-Nelson, Dhruva Tirumala, Hubert Soyer, Joel Leibo, Remi Munos, Charles Blundell, Dharshan Kumaran, and Matt Botvinick. Learning to reinforcement learn. arXiv preprint
2016.
Huaxiu Yao, Ying Junzhou Huang, and Zhenhui Li. Hierarchically structured meta-learning. In International Conference on Machine Learning, pp. 7045-7054. PMLR, 2019.
Jin Zhang, Jianhao Wang, Hao Hu, Tong Chen, Yingfeng Chen, Changjie Fan, and Chongjie Zhang. Metacure: Meta reinforcement learning with empowermentddrven exploration. In International Conference on Machine Learning, pp. 12600-12610. PMIR, 2021.
Wenxuan Zhou, Lerrel Pinto, and Abhinav Gupta. Environment probing interaction policies. In International Conference Learning Representations, 2019. URL /RR[itttps///peeeee/eee
Luisa Zintgraf, Kyriacos Shiarlis, Maximilian Igl, Sebastian Schulze, Yarin Gal, Katja Hofmann, and Shimon Whiteson. Varibad: A very good method for bayes-adaptive deep via meta-
learning. In International Conference on Learning Representations, 2020. URL https:
//openeeview.ne///orrmm///

---

Under review as a conference paper at ICLR 2023
A ELBO DERIVATION
E BptMett) [logp(+*)]=
Bp In p(z,co,tt r+) ((,,
(2,
p(2,c,tt)
1pp In ((
(2,
In p(2,
2 Pp (
(2,
= EELOO
where the inequality holds due to the Jensen's ' inequality. The inner expectation term can be further decomposed as,
In c(2,c,tt)
E (a
((2,cTT.))
p(r ch
R p
((1c,, (c)
p(2|c) + In p(c)
Inp(r+ 12, + In <(1c,, Ta)
((cITa))
Inp(T + 1 ] - By(d))))
(L(g(2)e,Ta) Ta) p(2lc)) (L(g(c)Ta) |l p(c)),
which is essentially 4). In practice, we first apply the Gumbel-Softmax trick (Jang et al. 2017) to sample a ē to calculate the second term instead directly calculating the expectation, and then sample given č to calculate the first term using the reparameterization trick (Kingma & &Welling 2013). To calculate the last term, we decompose into,
(g(((()) p(c)) = H [a(clTa)] + lEttt F.)) [p(c)]:
where p(c) is the uniform prior of the cluster distribution, thus the last KL term is essentially the entropy of the posterior ((cTT:e) plus a constant value, which can be easily computed.
B ALGORITHMS
Algorithm 1: MILET: Meta-train Phase
Input: A set of meta-train tasks M drawn from p(M); Initialize buffer for CBVI training;
while not Done do Sample a task M₁- ~ M;
Collect exploration and exploitation episodes ++ = [d) by running Alg. 2 on M₁;
Insert + to V;
Compute intrinsic rewards rs. and ra:
Train Ty on 7 by maximizing Eq.
using PPO:
Train T. on 7 by maximizing Eq. using PPO;
Sample trajectory batch from V
Update 0, w using v by maximizing Eq.
end
Algorithm [ summarizes the training procedure of MILET. We use the on-policy algorithm PPO (Schulman et al. 2017) to train policies; and retain a buffer storing trajectories for the cluster-based variational inference training. Algorithm 2 is used to obtain trajectories given task. When it is executed in Algorithm during training, is sampled using the Gumbel-Softmax trick, otherwise using argmax when handing new tasks during testing.
12

---

Under review as a conference paper at ICLR 2023
Algorithm 2: MILET: Meta-test Phase
Input: Meta-test task drawn from p(M), number of exploitation episodes N;; for i 1,... H7 do
Obtain arg max, 2o (c(ch,,,--)):
// Use Gunbel-Softmax during training
Take action according to Te (arlst, go (0(2 p.c.,
Update hy with (8t, ae, rr, St41);
end
Initialize hi gr for e = 1, N,,,,.,. .N do
for 1, ..... H do
Obtain arg max, 9o (c(cha aar): // Use Gunbel-Softma during training
Take action according ne (arels go (= (1-1)
Update he with (st, ae 8(+1;)
end end
C ENVIRONMENTS
Ant-Goal. We create clusters by manipulating the goal positions. Given angle 0 [0, obtain the 2D coordinate of the goal position as (r cos 0, rin 0), where r is the radius and fixed to 2 in our experiments. To create clusters, sample 0 from 4 different normal distributions (0.25, 0.2²), (00.75, 0.2?), (11255 0.2²), (1.775 0.22). To sample a task, we first uniformly sample a normal distribution, and then sample from
Humanoid-De We create clusters by manipulating the goal directions. To sample a goal direction, first sample angle 0 E [0, 2). and then the goal direction is (cos 0, sin 0). To create clusters, we sample from 4 different normal distributions (0.25, 0.2?). (0.7755, 0.2?), (11255, 0.2?), (L.77, 0.2?). To sample task, we first uniformly sample normal distribution, and then sample a 0 from it.
Hopper- Rae---andd-Paaa- We create clusters by manipulating physical parameters the robot, caus- ing different transition functions. We have four sets of parameters in total, including body mass, damping on degrees of freedom, body inertia and geometry friction. To sample a task, we first uni- formly sample one of the four parameter sets, and then multiply parameters in it with multipliers sampled from (3,, 1.5²).
The parameter sets are the same as the hopper robot. We find small mul- tipliers cannot change transition functions a lot in this environment, thus we sample the multipliers from (6, 12).
Walker- Walker-Rand-
Ant-Goal-Parrtee The tasks are sampled in the same way as Ant-Goal. The goal reward is defined as,
r= bc
g11 + 1 | 11
otherwise
where is the current coordinate of the robot, g is the coordinate of the goal position, t is the threshold, whish is fixed our setting. Control cost and contact cost are also included in the final reward.
Humanoid-Din Dir-airreee--eeee....e The tasks are sampled in the same way as Humanoid-Dry The velocity reward is defined as
r= f=(0
(g.wv U. I|lll|ell 2t otherwise
where v is the velocity of the robot, g is the goal direction, is the threshold and fixed to 0.8 in our setting. Control cost and contact cost also included in the final reward.
Ant-Goal-U, For each task, we uniformly sample a U(00 2) and obtain the goal position crs rsi rin 0) with fixed to 2.
13

---

Under review as a conference paper at ICLR 2023
0.100
0.075
Value
0.005
-0.025
00.075
60
100
Frames
Figure 6: Visualization of y in the exploration policy.
D HYPER-PARAMETERS
We used the Py Torch framework Paszke et al. (2019) to implement our experiments. We use PPO with Humber loss update MILET. In the cluster-based variation inference, we set À to and Xp to 0.1. We update the target network (z/c) every 50 training epochs. For the decaying function used in exploration policy, we set (t) = 0.1 0.1 exp(-0.1(H - +)) and 70(t) -0.1 0.2 0.2exp(-0.1((
t)), which are visualized in Figure 6 We use the Adam optimizer (Kingma Ba 2014) for both policy learning and cluster-based variational inference training. For policy learning including Te and Tg. set the learning rate le-4, while 1e-3 for variational inference training. In the clustered reward function environments, A set to in MiLEt and all variational inference baselines. In the clustered state transition function environments, X, is set to be 1. The other parameters are set as default. The max episode length H set to 100 for every environment.
We provide implementation of MILET and task generation process in the supplementary material.
E ABLATION STUDY
We study the contribution of proposed components in MILET. Firstly, we disabled the exploration policy and directly train the exploitation policy maximize task rewards. Secondly, we removed the stacked GRU (S-GRU) structure and only keep single GRU to encode the trajectory Ta in the cluster-based variational inference. Lastly, we disable the consistency regularizers (CR) in the variational inference and directly optimize the ELBO. We use the same environments as Section
The final performance of all variants is shown in Table 2 Firstly, without the exploration policy, the performance decreases significantly. This variant can still unravel clustering structures of tasks, so it performs better than VariBAD in general, which uses a simpler variational inference component. This again proves the necessity of handling structured heterogeneity in task distribution. However, cannot effectively utilize local knowledge as the exploration does not aim to explore different cluster structures. Our exploration policy utilizes the cluster structure information and is trained to explore the most certain clusters for tasks, leading to better adaptation performance. Secondly, removing S- GRU makes harder utilize patterns with different levels of granularity the interaction history. This ability is essential to quickly identify cluster structures as cluster patterns are more common across tasks, without which the exploration is less effective, causing performance drop. Also, consistent variational inference is very important. Our designed consistency regularizers are helpful discover cluster structures following the nature of sequential data in meta-RL, improving the final performance of MILET.
VISUALIZATION OF STATE TRANSITION FUNCTIONS
We visualize the clusters clustered transition function experiments in Figure 7 visualize the state transition function as it is not simple tabular function and it is conditioned on agent's actions. Recall that we created the ground-truie clusters by applying different multipliers on the initial parameters of each of the parameter sets, and thus such multipliers well characterize
It hard to directly
14

---

Under review as a conference paper at ICLR 2023
Table 2: Ablation analysis of MILET.
Ant-Gioa
Humanoid-Dir
Hopper-Rand-e Rand-Params
VariBAD-G
-171.2.9.7
514.3_24.6
272.9111.3
153.619.4
exploration
-130.414.5
533.9_18. 18.4
278.319.6
164.21166.7
S-GRU
-173.7+14.2
554.2=12.7
284,4+11.3
211.7+9.8
CCR
-1.2.8+8.7
528. 35.6
301.6+124
219.4+11.1
MILET
-125.425.1
277.7=28.0
312.2+18.1
232.6:115.4
the state transition in each task. We then calculate the average multipliers of tasks in each inferred cluster [
Ideally, one cluster of tasks should only have one of multipliers larger than 1, as is how we created different task clusters. From this visualization results, we can clearly observe that MILET successfully identifies the difference of clusters.
2.5
0.0
multipliers
2.5
0.0
2.5
0.0
body, mass
body jnertia
damping
geom on
bady mass
body. damping
geam_friction
#Parameter
#Parameter
(a) Hopper- Rand-Ranaaaa Params.
(b) Walker Malker-Rand-ae Hand-Params.
Figure Clusters in state transition functions. Each color represents identified cluster by MiLET. Four parameter sets are split by dashed lines.
G VISUALIZATION ON ANT-U
We visualize the task clusters identified by MILET-8 and MILET-10 on the testing tasks in the Ant- environment. We can observe that tasks are split into smaller groups by these two new variants, and their cluster assignments modeling causes ineffective exploration and sharing of wrong knowledge across tasks, leading to the degeneration of final performance in Table
mixed at the boundary of adjacent clusters. Such inaccurate cluster
MiLEt-B
MiLEt-10
Figure 8: Cluster assignments produced by MILET-8 and MILET-10 in the Ant-U environment.
"Multipliers with value 0 the corresponding initial parameters O.
15

---

Under review as a conference paper at ICLR 2023
H DISCUSSION OF CLUSTER-AWARE EXPLORATION
Our exploration policy distinguishes MILET from other task inference methods in utilizing the struc- of the task distribution. VariBAD Zintgraf al. 2020 performs exploration function of task uncertainty by maximizing task rewards. MetaCURE (Zhang et al. 2021) and DREAM (Liu al. 202T) propose explore by maximizing the mutual information between inferred task em- beddings and pre-defined task descriptions. These exploration methods are unaware of structures of the task distribution, thus less effective in exploring coarser but useful information, i.e., clusters. Our cluster-aware exploration is designed to quickly reduce the uncertainty in cluster inference. As shown in Figure 7a the agent explores on the map scale find the most suitable structure.
To better understand the implication of our cluster-aware exploration, we compare the clustering quality (measured by NMI score) at the end of the first episode MILET and its variant without exploration policy in Table Similar to VariBAD, this variant performs exploration by considering task uncertainty. Our exploration policy with explicit clustering objective obtains better clustering quality, which builds foundation for refining task inference, resulting in better final performance in Table
Table 3: NMI score of different exploration methods.
Ant-Gial
Humanoid-Di Dir
Rand- Rand---aandd
eexploration
0.466
0.485
0.517
0.327
MiLEr
0.615
0.639
0.588
0.765
16


Reviews:
[{"review_id": "wZc--fHncf", "reviewer": "Reviewer_XnEo", "summary": "", "questions": "", "limitations": "", "rating": 5, "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": "This work proposes a method, MILET, that leverages an assumption of heterogeneity within the task distribution. More specifically, the task distribution is modelled as a mixture distribution. The objective of MILET essentially adapts VariBAD (Zintgraf et al., 2020), decomposing the belief as a mixture id and a task latent conditioned on the mixture id. The sampling of beliefs for the task decoder is restricted to the inferred posterior mixture id, operationalizing the intuition of eliminating the transfer between unrelated tasks. Two additional regularization terms are also proposed. MILET also separately maintains exploration and exploitation agents, similar to DREAM (Liu et al., 2021). Experiments are run on MuJoCo-based meta-RL domains in which the task variation manifests in the reward function (2 domains) or dynamics (2 domains). Overall, MILET is shown to outperform meta-RL methods from prior works that do not assume a mixture task distribution.", "summary_of_the_review": "I currently recommend borderline reject due to the issues of clarity regarding the problem statement, limited technical novelty, and moderate empirical significance.", "strength_and_weaknesses": "Strengths\n\nThe authors study a hitherto neglected property of realistic task distributions. The proposed method that leverages this property is sensible and includes new terms that address the specific challenges arising from inferring mixture ids. The empirical evaluation involved modifying existing benchmark tasks to reflect this property.\n\nWeaknesses\n\nThe newly proposed meta-RL setting is not clearly distinguished from that of prior work. The authors mention that prior meta-RL assumes tasks are IID, but that seems to need to hold for this work as well. Assuming latent variables underlying the task distribution is well and good, but does the marginal task distribution need to exhibit clusters, or can it be relatively uniform? Put another way, If MILET can improve over prior methods for densely sampled task distributions, then perhaps the real benefit of MILET comes from leveraging the sufficiency of nearby tasks for transfer (a local property), as opposed to heterogeneity (a global property)? The novelty of the proposed method is limited as its backbone is essentially VariBAD (for the evidence lower bound) and MetaCURE/DREAM (for decoupled exploration and exploitation) endowed with the mixture task distribution assumption.\n\nMinor\n\n- I'm not a fan of the cooking analogy in the introduction. There is probably a lot of commonality in cooking Chinese and Indian dishes, which undercuts the point you're trying to make.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "b8rEE0ZUkl", "reviewer": "Reviewer_nQz2", "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": "This paper introduces a method for explicit task clustering as an inductive bias for Meta-RL algorithms operating on a task distribution. The key assumption is heterogeneity within that task distribution, i.e. that the i.i.d. assumption of draws from the task distribution is unrealistic in real Meta-RL setting and that knowledge can be more readily transferred by explicitly modelling shared structure. To that end, the authors propose a hierarchical probabilistic model, along with an inference procedure that first infers cluster ids and then the task within that cluster. In addition, the method further shows a training procedure distinguishing between an exploratory policy for the first and an exploiting policy for subsequent episodes and introduced various intrinsic rewards designed to lead to stable inference. The method shows promising results, although tasks are of limited complexity and clusters correspond to low-level information (Goal region).", "summary_of_the_review": "Overall a good submission with solid methodological improvements, albeit with somewhat limited complexity in its experimental section. In its current stage it is difficult to asses whether the method along with various introduced tricks scales to more complex scenarios, making it difficult to estimate the likely significance of this work in inspiriting future approaches. I'd consider raising my score if the authors improve the work in that regard during the rebuttal. Nevertheless a good submission on a challenging problem that would not be out of place as an ICLR publication. ", "strength_and_weaknesses": "`Strengths:`\n- Task clustering in Multi-task and Meta-RL is an important yet challenging problem and various attempts at designing such algorithms that make use of clustering have been unsuccessful. The authors correctly identified this as an important and timely problem. \n- Clarity in the algorithmic derivation and motivation of the work (Several minor suggestions for improvement are provided).\n- Results on the provided experiments are strong.\n\n`Weaknesses:`\n- The provided experiments are relatively simple and hand-designed, cluster information is relatively local and of limited abstraction. At this point I would be unsure whether the technique would be capable of clustering entire tasks from e.g. the MetaWorld task suite based on the more abstract problem being solved (e.g. various pushing/pull tasks, different tasks involving the same object (e.g. buttons, coffee mugs)).\n- The method does not handle unknown number of clusters. \n\n`Requests/Questions:`\n- Could the authors kindly provide a cluster visualisation of the clustered transition function experiments? \n- Table 1: At what point do increasing number of clusters reduce performance? ", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "4P3d0AS5i2", "reviewer": "Reviewer_cxKJ", "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": "The authors propose an inference based meta-learning approach, similar to Pearl or Varibad, but with a hierarchical latent task embedding modelled as GMM. Additionally, they propose to learn two policies, one optimised for exploration, one for exploitation. The algorithm is evaluated on MuJoCo tasks with variable goal locations which are sampled from clusters. They refer to this as \"structured heterogeneity\". ", "summary_of_the_review": "The paper proposes an interesting algorithm which seems to perform well in practise on tasks with \"structured heterogeneity\". \nMy main concern is with the discussion of the results and the justification of the algorithm. In particular, I am not convinced that the superior performance is due to the latent structure (as they authors claim). Additional discussion and, in particular, ablation studies would hence strenghten the paper. \n\n\n", "strength_and_weaknesses": "# Strenghts\n\n* The paper and problem setting is well motivated.\n* The related work section is well written, but could be more extensive.\n* The proposed ideas are interesting. \n\nI personally find especially the way the exploration policy is learned quite interesting. For one, splitting exploration and exploitation policy might be a useful inductive bias in itself as it might provide a simple but effective way of balancing both aspects -- as opposed to trying to do reward engineering to learn one policy fulfilling both roles.\nSecondly, the cluster entropy reward might be a very effective proxy reward for exploration success.\n\n# Weaknesses\n\nWhile I like the algorithm and can believe that it outperforms the baseline on the evaluated tasks, I am unsure about some of the justifications made for design choices and explanations surrounding the experimental results. I will discuss several instances of this below.\n\n## Where does the advantage come from: latent hierarchy or exploration policy?\n\nUnfortunately, the paper does very little to disentangle whether the success is coming from the clustering of the latent space or the separation between exploration and exploitation policy. Instead, it attributes most of its success to the clustering in latent space, claiming that this provides advantages in terms of knowledge sharing between tasks - which I'm not yet convinced by. In particular, I am not sure why VariBAD or Pearl would fail at effectively sharing that knowledge as their formulation seems general enough to do so.\nAnd while MILET provides a hierarchical inductive bias in the latent space, the success of non-hierarchical VAEs has, imo, shown that these types of flat latent spaces are usually flexible enough.  \n\n## Consistency\n\nIn section 4.1 the paper discusses a problem with \"consistency\". However, it is nowhere defined what exactly is meant by consistency. In particular, I am not sure how this can be a contradiction to the update to the agents belief -- which should become more certain over time. \n\nI assume that the authors mean by \"non-consistent\" if q(c|tau) 'jumps' from one time-step to the next in a way that the two KLs have very little overlap -- but either way, this should be made precise.  \n\n## Experimental Results\n\nI find it quite surprising that MILET's _exploration_ policy outperforms the other methods by such a large margin -- this seems to me to be more concistent with the hypothesis that the split between exploration and exploitation policy is the real innnovation here and not the GMM in the latent space?\n\nFor the sparse reward setup, I am not sure I understand the following sentence:\n\n> With sparse reward, it becomes more important for the agent to leverage knowledge across related tasks, as the feedback within a single task is insufficient for policy learning.\n\nGiven that the reward it sparse, how can knowledge sharing help: Before we find the goal, we cannot know which cluster it is in and after we find the goal, this becomes irrelevant as we've already found the goal. Again, imo this could be very well explained by the superior exploration policy and not the latent structure. ", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "Not applicable"}, "extra_text": {}}, {"review_id": "3Qq9YWIr69", "reviewer": "Reviewer_XdEV", "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": "This work proposes a new exploration objective in meta-reinforcement learning to discover clusters of related tasks. The claim is that tasks may be heterogeneous, where knowledge may only be transferable to other tasks within the same cluster, so exploring to identify these clusters could find the useful transferable information.", "summary_of_the_review": "Overall, I generally like the proposed method ideas. Exploration via clustering is original and can clearly be helpful. However, the framing and rhetoric of this work is a severe weakness, so I currently cannot recommend acceptance. I look forward to the author response.", "strength_and_weaknesses": " ## Strengths\n\n- The idea to create an exploration objective based on identifying clusters of related tasks is interesting and completely novel.\n- The experimental results are strong, although the experimental setup is imperfect, as I comment upon below.\n- By and large, the presentation is quite clear and sufficiently technically precise.\n\n## Weaknesses\n\nA major weakness of this work is that the rhetorical claims and explanations are often unsupported, and in several places, clearly false. Many claims appear to be speculation, while they seem to be presented as fact. Below, I include several examples:\n- This includes the primary motivation of this work: i.e., that tasks may be heterogeneous, where only some tasks may contain transferable information for other tasks, so it is crucial to identify clusters of transferable information. While identifying these clusters clearly holds some value, it is somewhat misleading, as identifying such clusters is _not_ necessary (i.e., a Bayes-adaptive optimal policy may achieve optimal returns, without any explicit knowledge of such clusters). Claims in the experiments like \"such exploration is not effective enough since it is not able to recover underlying clustered task structures\" are therefore misleading. Further, such claims are not even supported by the experiments themselves: results in the Appendix, which show that ablating the exploration policy from MILET continues to outperform VariBAD raise questions about what the primary source of empirical improvements is. As MILET outperforms baselines without even exploring to identify clusters, it's unclear if the gains are just from superior neural network architectures or not -- notably the baselines do not use the same architectures as MILET, and sharing the same architecture could help improve the experiments. Experiments on Ant-U without heterogeneous clusters, which are quite interesting, also suggest that it's not about finding transferability between tasks, and clustering can simply act as a useful auxiliary task / exploration objective. Finally, the claim that some tasks contain no transferable information to other tasks does not even appear to be true in the evaluated tasks or provided examples. In the example where clusters might be goals toward the top or the bottom, there are still shared motor skills. To be clear, I think the idea of clustering is still quite interesting, but proposing a story that is not supported is concerning.\n- The experiments also speculate about the cause of reduced performance from MetaCURE, but present this as fact: \"as in our setting all test tasks are new, the task id becomes meaningless in meta test. MetaCURE thus fails to generalize.\" MetaCURE does not use the task ID at meta-test time, so this is false.\n\nIt should be noted that exploration based on clustering can result in failure modes and is not always applicable. For example, in a distribution of tasks, where tasks can be clearly clustered into task-irrelevant but highly prominent features can lead to unhelpful exploration. In this case, the proposed reward bonuses may mislead exploration from finding actually useful information. Inclusion of such discussion could improve this work. As a related note, the tasks from the experiments do not in general require significant exploration, as much can be inferred from a single timestep reward from relatively random actions (although the exact reward details are not included for many of the tasks). This makes it somewhat challenging to evaluate the proposed exploration bonuses.\n\n## Minor Comments\n- Calling Ant-Goal-Sparse a \"sparse\" reward environment seems misleading if the rewards kick in even from the initial state and are not binary.\n- Why are there only 5 colors in Figure 5, if it is MILET-6?\n- Is the decomposition of the entropy H(q(c | tau)) into a telescoping series to create per-timestep rewards based on the similar decomposition from DREAM?\n\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Decision: Reject