Original: U NDERSTANDING THE C OMPLEXITY G AINS OF C ONTEXTUAL M ULTI - TASK RL WITH  C URRICULA
Revised: U NDERSTANDING THE C OMPLEXITY G AINS OF R EFORMULATING S INGLE -T ASK RL WITH A C URRICULUM

Original: Anonymous authors
Revised: Anonymous authors

Original: Paper under double-blind review
Revised: Paper under double-blind review

Original: A BSTRACT
Revised: A BSTRACT

Original: Reinforcement learning (RL) problems can be challenging without well-shaped rewards.
Revised: Reinforcement learning (RL) problems can be challenging without well-shaped rewards.

Original: Prior work on provably efficient RL methods generally proposes to address this issue with dedicated exploration strategies, such as novelty-based bonuses.
Revised: Prior work on provably efficient RL methods generally proposes to address this issue with dedicated exploration strategies.

Original: However, another way to tackle this challenge is to reformulate it as a multi-task RL problem, where the task space contains not only the challenging task of interest but also easier tasks that implicitly function as a curriculum.
Revised: However, another way to tackle this challenge is to reformulate it as a multi-task RL problem, where the task space contains not only the challenging task of interest but also easier tasks that implicitly function as a curriculum.

Original: Such a reformulation opens up the possibility of running existing multi-task RL methods as a more efficient alternative to solving a single challenging task from scratch.
Revised: Such a reformulation opens up the possibility of running existing multi-task RL methods as a more efficient alternative to solving a single challenging task from scratch.

Original: In this work, we provide a theoretical framework that reformulates a single-task RL problem as a multi-task RL problem defined by a curriculum.
Revised: In this work, we provide a theoretical framework that reformulates a single-task RL problem as a multi-task RL problem defined by a curriculum.

Original: Under mild regularity conditions on the curriculum, we show that sequentially solving each task in the multi-task RL problem is more computationally efficient than solving the original single-task problem, without any explicit exploration bonuses or other exploration strategies.
Revised: Under mild regularity conditions on the curriculum, we show that sequentially solving each task in the multi-task RL problem is more computationally efficient than solving the original single-task problem, without any explicit exploration bonuses or other exploration strategies.

Original: We also show that our theoretical insights can be translated into an effective practical learning algorithm that can accelerate curriculum learning on robotic goal-reaching tasks.
Revised: We also show that our theoretical insights can be translated into an effective practical learning algorithm that can accelerate curriculum learning on robotic goal-reaching tasks.

Original:  I NTRODUCTION
Revised:  I NTRODUCTION

Original: Reinforcement learning (RL) provides an appealing and simple way to formulate control and decisionmaking problems in terms of reward functions that specify what an agent should do, and then automatically train policies to learn how to do it.
Revised: Reinforcement learning (RL) provides an appealing and simple way to formulate control and decisionmaking problems in terms of reward functions that specify what an agent should do, and then automatically train policies to learn how to do it.

Original: However, in practice the specification of the reward function requires great care: if the reward function is well-shaped, then learning can be fast and effective, but if rewards are delayed, sparse, or can only be achieved after extensive explorations, RL problems can be exceptionally difficult (Kakade and Langford, 2002; Andrychowicz et al., 2017;
Revised: However, in practice the specification of the reward function requires great care: if the reward function is well-shaped, then learning can be fast and effective, but if rewards are delayed, sparse, or can only be achieved after extensive explorations, RL problems can be exceptionally difficult (Kakade and Langford, 2002; Andrychowicz et al., 2017;

Original: Agarwal et al., 2019).
Revised: Agarwal et al., 2019).

Original: This challenge is often overcome with either reward shaping (Ng et al., 1999;
Revised: This challenge is often overcome with either reward shaping (Ng et al., 1999;

Original: Andrychowicz et al., 2017; 2020) or dedicated exploration methods (Tang et al., 2017; Stadie et al., 2015; Bellemare et al., 2016; Burda et al., 2018), but reward shaping can bias the solution away from optimal behavior, while even the best exploration methods, in general, may require covering the entire state space before discovering high-reward regions.
Revised: Andrychowicz et al., 2017; 2020; Gupta et al., 2022) or dedicated exploration methods (Tang et al., 2017; Stadie et al., 2015; Bellemare et al., 2016; Burda et al., 2018), but reward shaping can bias the solution away from optimal behavior, while even the best exploration methods, in general, may require covering the entire state space before discovering high-reward regions.

Original: On the other hand, a number of recent works have proposed multi-task learning methods in RL that involve learning contextual policies that simultaneously represent solutions to an entire space of tasks, such as policies that reach any potential goal (Fu et al., 2018;
Revised: On the other hand, a number of recent works have proposed multi-task learning methods in RL that involve learning contextual policies that simultaneously represent solutions to an entire space of tasks, such as policies that reach any potential goal (Fu et al., 2018;

Original: Eysenbach et al., 2020b; Fujita et al., 2020; Zhai et al., 2022), policies conditioned on language commands (Nair et al., 2022), or even policies conditioned on the parameters of parametric reward functions (Kulkarni et al., 2016;
Revised: Eysenbach et al., 2020b; Fujita et al., 2020; Zhai et al., 2022), policies conditioned on language commands (Nair et al., 2022), or even policies conditioned on the parameters of parametric reward functions (Kulkarni et al., 2016;

Original: Siriwardhana et al., 2019; Eysenbach et al., 2020a; Yu et al., 2020b).
Revised: Siriwardhana et al., 2019; Eysenbach et al., 2020a; Yu et al., 2020b).

Original: While such methods are often not motivated directly from the standpoint of handling challenging exploration scenarios, but rather directly aim to acquire policies that can perform all tasks in the task space, these multi-task formulations often present a more tractable learning problem than acquiring a solution to a single challenging task in the task space (e.g., the hardest goal, or the most complex language command).
Revised: While such methods are often not motivated directly from the standpoint of handling challenging exploration scenarios, but rather directly aim to acquire policies that can perform all tasks in the task space, these multi-task formulations often present a more tractable learning problem than acquiring a solution to a single challenging task in the task space (e.g., the hardest goal, or the most complex language command).

Original: We pose the following question: can we construct a multi-task RL problem with contextual policies that is easier than solving a single-task RL problem?
Revised: We pose the following question: can we construct a multi-task RL problem with contextual policies that is easier than solving a single-task RL problem from scratch?

Original: In this work, we answer this question affirmatively by analyzing the sample complexity of a class of curriculum learning methods based on this principle. [Figure]
Revised: In this work, we answer this question affirmatively by analyzing the sample complexity of a class of curriculum learning methods.

Original: To build the intuition for how reformulating a single-task problem into a multi-task problem enables efficient learning, consider the setting where the optimal state visitation distributions  of two different contexts ω, ω ′ are “similar”, and our goal is to learn the optimal policy π ⋆ω ′ w.r.t.
Revised: To build the intuition for how reformulating a single-task problem into a multi-task problem enables efficient learning, consider the setting where the optimal state visitation distributions d π ⋆ω µ , d π ⋆ω ′ µ of two different contexts ω, ω ′ are “similar”, and our goal is to learn the optimal policy π ⋆ω ′ w.r.t.

Original: ω ′ .
Revised: ω ′ .

Original: Suppose we have learned the optimal policy π ⋆ω , we can facilitate learning π ⋆ω ′ by: (1) using π ⋆ω as an initialization and (2) setting a new initial state distribution µ ′ =
Revised: Suppose we have learned the optimal policy π ⋆ω , we can facilitate learning π ⋆ ω ′ by: (1) using π ⋆ω as an initialization and (2) setting a new initial state distribution µ ′ =

Original: βd π ⋆ω µ + (1 − β µ ) by mixing d π ⋆ω µ , the optimal state visitation distribution of π ⋆ω , and µ , the initial distribution of interest.
Revised: βd π ⋆ω µ + (1 − β µ ) by mixing d π ⋆ω µ , the optimal state visitation distribution of π ⋆ω , and µ , the initial distribution of interest.

Original: Using π ⋆ω asan initialization for learning π ⋆ω ′ facilitates the learning process as it guarantees the initialization is within some neighborhood of the optimality.
Revised: Using π ⋆ω as initialization for learning π ⋆ω ′ facilitates the learning process as it guarantees the initialization is within some neighborhood of the optimality.

Original: Using the new initial distribution µ ′ =
Revised: Setting the new initial distribution µ ′ =

Original: βd π ⋆ω µ + (1 − β ) µ also facilitates the learning of π ⋆ ω ′ as it reduces the density mismatch ratio between (cid:13)(cid:13)(cid:13)
Revised: βd π ⋆ω µ + (1 − β ) µ also facilitates the learning of π ⋆ω ′ as it reduces the density mismatch ratio between (cid:13)(cid:13)(cid:13)

Original: d π ⋆ω ′ µ ′ /µ ′
Revised: d π ⋆ω ′ µ ′ /µ ′

Original: (cid:13)(cid:13)(cid:13) ∞ , a key quantity that influences sample complexity.
Revised: (cid:13)(cid:13)(cid:13) ∞ , a key quantity that influences sample complexity.

Original: The new distribution µ ′ =
Revised: The new distribution µ ′ =

Original: βd π ⋆ω µ + (1 − β ) µ for learning π ⋆ω ′ could be understood as “rolling in d π ⋆ω µ with probability β ”.
Revised: βd π ⋆ω µ + (1 − β ) µ for learning π ⋆ω ′ could be understood as “rolling in d π ⋆ω µ with probability β ”.

Original: we refer  this approach, which consists of rolling in a similar optimal state visitation distribution of another context , as R OLLIN .
Revised: We refer to this approach, which consists of rolling in a similar optimal state visitation distribution of another context , as R OLLIN .

Original:  We provide an illustration of R OLLIN in Figure 1, which represents the basic idea of R OLLIN .
Revised: [Figure] We illustrate the intuition of R OLLIN in Figure 1.

Original: Our paper analyzes this algorithm and proves that it can attain polynomial-time sample complexity, without any explicit exploration bonuses

Original: We formalize this intuition in the contextual MDP setting, where we assume each MDP M ω is uniquely defined by a context ω in the context space W ⊂ R n , and we are given a curriculum { ω k } Kk =0 , with the last MDP M ω K being the MDP of interest.
Revised: More specifically, we adopt the contextual MDP formulation, where we assume each MDP M ω is uniquely defined by a context ω in the context space W ⊂ R n , and we are given a curriculum { ω k } Kk =0 , with the last MDP M ω K being the MDP of interest.

Original: To show our main results, we only require a Lipschitz continuity assumption on r ω w.r.t.
Revised: To show our main results, we only require a Lipschitz continuity assumption on r ω w.r.t.

Original: ω and some mild regularity conditions on the curriculum { ω k } Kk =0 .
Revised: ω and some mild regularity conditions on the curriculum { ω k } Kk =0 .

Original: With the aforementioned assumptions, we show that learning π ⋆K by recursively rolling in with a near optimal policy for ω k to construct the initial distribution µ k + for the next context ω k +1 , is provably more efficient than learning π ⋆ω K from scratch.
Revised: We show that learning π ⋆K by recursively rolling in with a nearoptimal policy for ω k to construct the initial distribution µ k +1 for the next context ω k +1 , is provably more efficient than learning π ⋆ω K from scratch.

Original: In particular, we show that when an appropriate sequence of contexts is selected, we can reduce the iteration and sample complexity bounds of entropy regularized softmax policy gradient (with an inexact stochastic estimation of the gradient) from an original exponential dependency on the state space size, as suggested by Ding et al.
Revised: In particular, we show that when an appropriate sequence of contexts is selected, we can reduce the iteration and sample complexity bounds of entropy-regularized softmax policy gradient (with an inexact stochastic estimation of the gradient) from an original exponential dependency on the state space size, as suggested by Ding et al.

Original: (2021), to a polynomial dependency.
Revised: (2021), to a polynomial dependency.

Original: We also prescribe a practical implementation of R OLLIN .
Revised: We also prescribe a practical implementation of R OLLIN .

Original: In summary, our contributionsof this work can be stated as follows.
Revised: In summary, our contributions can be stated as follows.

Original: We first provide a theoretical method (R OLLIN ) that facilitates single-task policy learning by recasting it as a multi-task problem under entropy regularized softmax policy gradient (PG), which reduces the exponential complexity bound of the entropy regularized PG to a polynomial dependency on S .
Revised: We first provide a theoretical method (R OLLIN ) that facilitates single-task policy learning by recasting it as a multi-task problem under entropy-regularized softmax policy gradient (PG), which reduces the exponential complexity bound of the entropy-regularized PG to a polynomial dependency on S .

Original: Lastly, we also provide a deep RL implementation of R OLLIN and demonstrate adding R OLLIN improves performance in simulated benchmark tasks in various settings, including goal relabeling (Andrychowicz et al., 2017) and Go-Explore (Ecoffet et al., 2019).
Revised: Last but not least, we also provide a deep RL implementation of R OLLIN and demonstrate adding R OLLIN improves performance in simulated goal-reaching tasks with an oracle curriculum and a non-oracle curriculum learned from MEGA (Pitis et al., 2020), as well as several standard Mujoco locomotion tasks inspired by meta RL (Finn et al., 2017).

Original:  R ELATED W ORK
Revised:  R ELATED W ORK

Original: Convergence of policy gradient methods.
Revised: Convergence of policy gradient methods.

Original: Theoretical analysis of policy gradient methods has a long history (Williams, 1992; Sutton et al., 1999; Konda and Tsitsiklis, 1999; Kakade and Langford, 2002; Peters and Schaal, 2008).
Revised: Theoretical analysis of policy gradient methods has a long history (Williams, 1992; Sutton et al., 1999; Konda and Tsitsiklis, 1999; Kakade and Langford, 2002; Peters and Schaal, 2008).

Original: Motivated by the recent empirical success (Schulman et al., 2015; 2017) in policy gradient (PG) methods, the theory community has extensively studied the convergence of PG in various settings (Fazel et al., 2018; Agarwal et al., 2021; 2020; Bhandari and Russo, 2019; Mei et al., 2020; Zhang et al., 2020b; Agarwal et al., 2020; Zhang et al., 2020a; 
Revised: Motivated by the recent empirical success (Schulman et al., 2015; 2017) in policy gradient (PG) methods, the theory community has extensively studied the convergence of PG in various settings (Fazel et al., 2018; Agarwal et al., 2021; 2020; Bhandari and Russo, 2019; Mei et al., 2020; Zhang et al., 2020b; Agarwal et al., 2020; Zhang et al., 2020a; Li et al., 2021;

Original: Li et al., 2021; Cen et al., 2021; Ding et al., 2021; Yuan et al., 2022; Moskovitz et al., 2022).
Revised: Cen et al., 2021; Ding et al., 2021; Yuan et al., 2022; Moskovitz et al., 2022).

Original: Agarwal et al.
Revised: Agarwal et al.

Original: (2021) established the asymptotic global convergence of policy gradient under different policy parameterizations.
Revised: (2021) established the asymptotic global convergence of policy gradient under different policy parameterizations.

Original: We extend the result of entropy regularized PG with stochastic gradient (Ding et al., 2021) to the contextual MDP setting.
Revised: We extend the result of entropy regularized PG with stochastic gradient (Ding et al., 2021) to the contextual MDP setting.

Original: In particular, our contextual MDP setting reduces the exponential state space dependency w.r.t.
Revised: In particular, our contextual MDP setting reduces the exponential state space dependency w.r.t.

Original: the iteration number and per iteration sample complexity suggested by Ding et al.
Revised: the iteration number and per iteration sample complexity suggested by Ding et al.

Original: (2021) to a polynomial dependency, when we have access to an appropriate curriculum.
Revised: (2021) to a polynomial dependency.

Original: We shall also clarify that there are many existing convergence analysis on other variants of PG that produce an iteration number that do not suffer from an exponential state space dependency (Agarwal et al., 2021; Mei et al., 2020), but they assume access to the exact gradient during each update of PG, while we assume a stochastic estimation of the gradient, which is arguably more practical.
Revised: We shall also clarify that there is much existing convergence analysis on other variants of PG that produce an iteration number that does not suffer from an exponential state space dependency (Agarwal et al., 2021; Mei et al., 2020), but they assume access to the exact gradient during each update of PG, while we assume a stochastic estimation of the gradient, which is arguably more practical.

Original: Contextual MDPs.
Revised: Contextual MDPs.

Original: Contextual MDPs (or MDPs with side information) have been studied extensively in the theoretical RL literature (Abbasi-Yadkori and Neu, 2014; 
Revised: Contextual MDPs (or MDPs with side information) have been studied extensively in the theoretical RL literature (Abbasi-Yadkori and Neu, 2014; Hallak et al., 2015; Dann et al.,

Original: Hallak et al., 2015; Dann et al., 2019; Jiang et al., 2017; Modi et al., 2018; Sun et al., 2019;
Revised: ; Jiang et al., 2017; Modi et al., 2018; Sun et al., 2019;

Original: Dann et al., 2019; Modi et al., 2020).
Revised: Dann et al., 2019; Modi et al., 2020).

Original: We analyze the iteration complexity and sample complexity of (stochastic) policy
Revised: We analyze the iteration complexity and sample complexity of (stochastic) policy

Original: gradient methods, which is distinct from these prior works that mainly focus on regret bounds (Abbasi-Yadkori and
Revised: gradient methods, which is distinct from these prior works that mainly focus on regret bounds (Abbasi-Yadkori and

Original: Neu, 2014; Hallak et al., 2015; Dann et al., 2019) and PAC bounds (Jiang et al., 2017; Modi et al., 2018; Sun et al., 2019; Dann et al., 2019; Modi et al., 2020).
Revised: Neu, 2014; Hallak et al., 2015; Dann et al., 2019) and PAC bounds (Jiang et al., 2017; Modi et al., 2018; Sun et al., 2019; Dann et al., 2019; Modi et al., 2020).

Original: Several works assumed linear transition kernel and reward model (or generalized linear model (Abbasi-Yadkori and Neu, 2014)) with respect to the context (Abbasi-Yadkori and Neu, 2014;
Revised: Several works assumed linear transition kernel and reward model (or generalized linear model (Abbasi-Yadkori and Neu, 2014)) with respect to the context (Abbasi-Yadkori and Neu, 2014;

Original: Modi et al., 2018; Dann et al., 2019; Modi et al., 2020; Belogolovsky et al., 2021).
Revised: Modi et al., 2018; Dann et al., 2019; Modi et al., 2020; Belogolovsky et al., 2021).

Original: These assumptions share similarity to our assumptions — we have a weaker Lipschitz continuity assumption with respect to the context space (since linear implies
Revised: These assumptions share similarity to our assumptions — we have a weaker Lipschitz continuity assumption with respect to the context space (since linear implies

Original: Lipschitz) on the reward function and a stronger shared transition kernel assumption.
Revised: Lipschitz) on the reward function and a stronger shared transition kernel assumption.

Original: Exploration.
Revised: Exploration.

Original: A number of prior works have shown that one can reduce the complexity of learning an optimal policy with effective exploration methods (Azar et al., 2017; Jin et al., 2018; Du et al., 2019; Misra et al., 2020; Agarwal et al., 2020; Zhang et al., 2020d).
Revised: A number of prior works have shown that one can reduce the complexity of learning an optimal policy with effective exploration methods (Azar et al., 2017; Jin et al., 2018; Du et al., 2019; Misra et al., 2020; Agarwal et al., 2020; Zhang et al., 2020d).

Original: The computational efficiency suggested by our work is different from some of the aforementioned prior methods that rely on adding exploration bonus for policy cover (Azar et al., 2017; Jin et al., 2018; Agarwal et al., 2020;
Revised: The computational efficiency suggested by our work is different from some of the aforementioned prior methods that rely on adding exploration bonus for policy cover (Azar et al., 2017; Jin et al., 2018; Agarwal et al., 2020;

Original: Zhang et al., 2020d), as we assume access to a “good” curriculum which ensures the optimal policy defined by the next context is not too different from the optimal policy of the current context.
Revised: Zhang et al., 2020d), as we assume access to a “good” curriculum which ensures the optimal policy defined by the next context is not too different from the optimal policy of the current context.

Original: The key intuition of our work is more related to the strategic exploration in the block MDP setting (Du et al., 2019; Misra et al., 2020).
Revised: The key intuition of our work is more related to the strategic exploration in the block MDP setting (Du et al., 2019; Misra et al., 2020).

Original: The block MDP setting assume the state space can be encoded into a feature space with a sequential block structure, and explores the next block in  feature space after having a policy that covers the previous blocks.
Revised: The block MDP setting assumes the state space can be encoded into a feature space with a sequential block structure and explores the next block in the feature space after having a policy that covers the previous blocks.

Original: Such sequential exploration procedure is similar to R OLLIN , as our method also considers solving the multi-task contextual MDPs in a sequential perspective, but R OLLIN focuses on learning the optimal policy of the next context once we have a near-optimal policy of the current context.
Revised: Such sequential exploration procedure is similar to R OLLIN , as our method also considers solving the multi-task contextual MDPs in a sequential perspective, but R OLLIN focuses on learning the optimal policy of the next context once we have a near-optimal policy of the current context.

Original: Curriculum learning in reinforcement learning.
Revised: Curriculum learning in reinforcement learning.

Original: Curriculum learning is a powerful idea that has been widely used in RLto address exploration challenges (Florensa et al., 2017; Kim and Choi, 2018; Omidshafiei et al., 2019; Ivanovic et al., 2019; Akkaya et al., 2019; Portelas et al., 2020;  Fang et al., 2020; Klink et al., 2020; Dennis et al., 2020; Parker-Holder et al.,  2022) (also see (Narvekar et al., 2020) for a detailed survey).
Revised: Curriculum learning is a powerful idea that has been widely used in RL (Florensa et al., 2017; Kim and Choi, 2018; Omidshafiei et al., 2019; Ivanovic et al., 2019; Akkaya et al., 2019; Portelas et al., 2020; Bassich et al., 2020; Fang et al., 2020; Klink et al., 2020; Dennis et al., 2020; Parker-Holder et al., 2022; Liu et al., 2022) (also see (Narvekar et al., 2020) for a detailed survey).

Original: Although curricula formed by well-designed reward functions (Vinyals et al., 2019; OpenAI, 2018; Berner et al., 2019; Ye et al., 2020; Zhai et al., 2022) are usually sufficient given enough domain knowledge, tackling problems with limited domain knowledge requires a more general approach where a suitable curriculum is automatically formed from a task space.
Revised: Although curricula formed by well-designed reward functions (Vinyals et al., 2019; OpenAI, 2018; Berner et al., 2019; Ye et al., 2020; Zhai et al., 2022) are usually sufficient given enough domain knowledge, tackling problems with limited domain knowledge requires a more general approach where a suitable curriculum is automatically formed from a task space.

Original: In the goal-conditioned reinforcement learning literature, this corresponds to automatic goal proposal mechanisms (Florensa et al., 2018; Warde-Farley et al., 2018; Sukhbaatar et al., 2018;
Revised: In the goal-conditioned reinforcement learning literature, this corresponds to automatic goal proposal mechanisms (Florensa et al., 2018; Warde-Farley et al., 2018; Sukhbaatar et al., 2018;

Original: Ren et al., 2019; Ecoffet et al., 2019; Hartikainen et al., 2019; Pitis et al., 2020; Zhang et al., 2020c; OpenAI et al., 2021; Zhang et al., 2021). 
Revised: Ren et al., 2019; Ecoffet et al., 2019; Hartikainen et al., 2019; Pitis et al., 2020; Zhang et al., 2020c; OpenAI et al., 2021; Zhang et al., 2021). The practical instantiation of this work is also similar to Bassich et al. Liu et al.

Revised: (2022), where a curriculum is adopted for learning a progression of a set of tasks.

Original: Learning conditional policies in multi-task RL.
Revised: Learning conditional policies in multi-task RL.

Original: Multi-task RL (Tanaka and Yamamura, 2003) approaches usually learn a task-conditioned policy that is shared across different tasks (Rusu et al., 2015; Rajeswaran et al., 2016;
Revised: Multi-task RL (Tanaka and Yamamura, 2003) approaches usually learn a task-conditioned policy that is shared across different tasks (Rusu et al., 2015; Rajeswaran et al., 2016;

Original: Andreas et al., 2017;  D’Eramo et al., 2020; Yu et al., 2020a; Ghosh et al., 2021; Kalashnikov et al., 2021).
Revised: Andreas et al., 2017; Finn et al., 2017; D’Eramo et al., 2020; Yu et al., 2020a; Ghosh et al., 2021; Kalashnikov et al., 2021).

Original: Compared to learning each task independently, joint training enjoys the sample efficiency benefits from sharing the learned experience across different tasks as long as the policies generalize well across tasks.
Revised: Compared to learning each task independently, joint training enjoys the sample efficiency benefits from sharing the learned experience across different tasks as long as the policies generalize well across tasks.

Original: To encourage generalization, it is often desirable to condition policies on low dimensional feature representations that are shared across different tasks instead (e.g., using variational auto-encoders (Nair et al., 2018; Pong et al., 2019; Nair et al., 2020) or variational information bottleneck (Goyal et al., 2019; 2020; Mendonca et al., 2021)).
Revised: To encourage generalization, it is often desirable to condition policies on low dimensional feature representations that are shared across different tasks instead (e.g., using variational auto-encoders (Nair et al., 2018; Pong et al., 2019; Nair et al., 2020) or variational information bottleneck (Goyal et al., 2019; 2020; Mendonca et al., 2021)).

Original: The idea of learning contextual policies has also been discussed in classical adaptive control literature (Sastry et al., 1990; Tao, 2003; Landau et al., 2011; Åström and Wittenmark, 2013; Goodwin and Sin, 2014).
Revised: The idea of learning contextual policies has also been discussed in classical adaptive control literature (Sastry et al., 1990; Tao, 2003; Landau et al., 2011; Åström and Wittenmark, 2013; Goodwin and Sin, 2014).

Original: Different from these prior works which have been mostly focusing on learning policies that can generalize across different tasks, our work focuses on how the near-optimal policy from a learned task could be used to help the learning of a similar task.
Revised: Different from these prior works which have been mostly focusing on learning policies that can generalize across different tasks, our work focuses on how the near-optimal policy from a learned task could be used to help the learning of a similar task.

Original:  P RELIMINARIES
Revised:  P RELIMINARIES

Original: We consider the contextual MDP setting, where a contextual MDP M W =
Revised: We consider the contextual MDP setting, where a contextual MDP M W =

Original: ( W , S , A , P , r ω , γ, ρ ) consists of a context space
Revised: ( W , S , A , P , r ω , γ, ρ ) consists of a context space

Original: W , a state space S , an action space A , a transition dynamic function P :
Revised: W , a state space S , an action space A , a transition dynamic function P :

Original: S ×
Revised: S ×

Original: A → P ( S ) (where P ( X ) denotes the set of all probability distributions over set X ), a context-conditioned reward function r :
Revised: A → P ( S ) (where P ( X ) denotes the set of all probability distributions over set X ), a context-conditioned reward function r :

Original: W × S × A → [0 , 1] , a discount factor γ ∈ (0 , 1] , and an initial state distribution of interest ρ .
Revised: W × S × A → [0 , 1] , a discount factor γ ∈ (0 , 1] , and an initial state distribution of interest ρ .

Original: For convenience we use S = |S| , A = |A| to denote the number of states and actions.
Revised: For convenience, we use S = |S| , A = |A| to denote the number of states and actions.

Original: While some contextual MDP formulations (Hallak et al., 2015) have context-conditioned transition dynamics and reward functions, we consider the setting where only the reward function can change across contexts.
Revised: While some contextual MDP formulations (Hallak et al., 2015) have context-conditioned transition dynamics and reward functions, we consider the setting where only the reward function can change across contexts.

Original: We denote r ω as the reward function conditioned on a fixed ω ∈ W and M ω =
Revised: We denote r ω as the reward function conditioned on a fixed ω ∈ W and M ω =

Original: ( S , A , P , r ω , γ, ρ ) as the MDP in- duced by such fixed reward function.
Revised: ( S , A , P , r ω , γ, ρ ) as the MDP in- duced by such fixed reward function.

Original: We use π ( a | s ) :
Revised: We use π ( a | s ) :

Original: S → P ( A ) to denote a policy and we adopt the softmax parameterization: π θ ( a | s ) =
Revised: S → P ( A ) to denote a policy and we adopt the softmax parameterization: π θ ( a | s ) =

Original: exp[ θ ( s,a )] (cid:80) a ′ exp[ θ ( s,a ′ )] , where θ : S × A (cid:55)→ R .
Revised: exp[ θ ( s,a )] (cid:80) a ′ exp[ θ ( s,a ′ )] , where θ : S × A (cid:55)→ R .

Original: We use d πρ ( s ) :
Revised: We use d πρ ( s ) :

Original: = (1 − γ ) (cid:80) ∞
Revised: = (1 − γ ) (cid:80) ∞

Original: t =

Revised: t =

Original: γ t P π ( s t = s | s 0 ∼ ρ ) to denote the discounted state visitation distribution and V πω :=
Revised: γ t P π ( s t = s | s 0 ∼ ρ ) to denote the discounted state visitation distribution and V πω :=

Original: E [ (cid:80) ∞ t =0 γ t r ω ( s t , a t )] + α H ( ρ, π ) to denote the entropy regularized discounted return on M ω , where H ( ρ, π ) :
Revised: E [ (cid:80) ∞ t =0 γ t r ω ( s t , a t )] + α H ( ρ, π ) to denote the entropy regularized discounted return on M ω , where H ( ρ, π ) :

Original: =
Revised: =

Original: E s 0 ∼ ρ,a h ∼ π ( ·| s h ) (cid:2)(cid:80)
Revised: E s 0 ∼ ρ,a h ∼ π ( ·| s h ) (cid:2)(cid:80)

Original: ∞
Revised: ∞

Original: h =0 − γ h log π ( a h | s h ) (cid:3) is the discounted entropy term.
Revised: h =0 − γ h log π ( a h | s h ) (cid:3) is the discounted entropy term.

Original: We use π ⋆ω := arg max π V πω to denote an optimal policy that maximizes the discounted return under M ω .
Revised: We use π ⋆ω := arg max π V πω to denote an optimal policy that maximizes the discounted return under M ω .

Original: We assume all the contextual reward functions are bounded within [0 , 1] : r ω ( s, a ) ∈ [0 , 1] , ∀ ω ∈ Ω , ∀ ( s, a ) ∈ S × A .
Revised: We assume all the contextual reward functions are bounded within [0 , 1] : r ω ( s, a ) ∈ [0 , 1] , ∀ ω ∈ Ω , ∀ ( s, a ) ∈ S × A .

Original: Similarly to previous analysis Agarwal et al.
Revised: Similarly to previous analysis Agarwal et al.

Original: (2021); Mei et al.
Revised: (2021); Mei et al.

Original: Ding et al.
Revised: Ding et al.

Original: (2021), we assume the initial distribution ρ for PG or stochastic PG satisfies ρ ( s ) > 0 , ∀ s ∈ S .
Revised: (2021), we assume the initial distribution ρ for PG or stochastic PG satisfies ρ ( s ) > 0 , ∀ s ∈ S .

Original: 3. ASSUMPTIONS
Revised: 3. A

Revised: [Equation]

Original: Given a curriculum { ω k } Kk =0 , where the last context ω K defines M ω K , the MDP of interest, our goal is to show that reformulating M ω K into a multi-task problem {M ω k } Kk =0 , enjoys better computational complexity and sample complexity than solving the single-task problem
Revised: Given a curriculum { ω k } Kk =0 , where the last context ω K defines M ω K , the MDP of interest, our goal is to show that reformulating M ω K into a multi-task problem {M ω k } Kk =0 , enjoys better computational complexity and sample complexity than solving the single-task problem

Original: M ω K from scratch.
Revised: M ω K from scratch.

Original: As we will show in Section 4, if we have a “good” curriculum { ω k } Kk =0 such that the optimal policies π ⋆ω k , π ⋆ω k +1 w.r.t.
Revised: As we will show in Section 4, if we have a “good” curriculum { ω k } Kk =0 such that the optimal policies π ⋆ω k , π ⋆ω k +1 w.r.t.

Original: two consecutive contexts ω k , ω k +1 are “close enough” to each other, then using an ε -optimal policy of ω k as an initialization allows us to directly start from the near optimal regime of ω k +1 , hence only requiring polynomial complexity to learn π ⋆ω k +1 .
Revised: two consecutive contexts ω k , ω k +1 are “close enough” to each other, then using an ε -optimal policy of ω k as an initialization allows us to directly start from the near optimal regime of ω k +1 , hence only requiring polynomial complexity to learn π ⋆ω k +1 .

Original: Formally, we use the two following assumptions to characterize the good properties of our curriculum.
Revised: Formally, we use the two following assumptions to characterize the good properties of our curriculum.

Original: Assumption 3.1 (Lipschitz reward in the context space)
Revised: Assumption 3.1 (Lipschitz reward in the context space)

Original: The reward function is Lipschitz continuous w.r.t.
Revised: The reward function is Lipschitz continuous w.r [Equation]

Original:  the context:
Revised: t. the context:

Original: [Equation]
Revised: [Equation]

Original: Assumption 3.2 (Similarity between two Consecutive Contexts)
Revised: Assumption 3.2 (Similarity between two Consecutive Contexts)

Original: We assume the given curriculum { ω k } Kk =

Revised: We assume the given curriculum { ω k } Kk =

Original: satisfies max 0 ≤ k ≤ K − 1 ∥ ω k +1 − ω k ∥ 2 ≤ O (cid:0) S − 2 (cid:1) , and  π ⋆ω 0 is also provided by the curriculum.
Revised: satisfies max 0 ≤ k ≤ K − 1 ∥ ω k +1 − ω k ∥ 2 ≤ O (cid:0) S − 2 (cid:1) , and we have access to a nearoptimal initialization θ (0)0 for learning π ⋆ω 0 (formally define in Section 4.2).

Original: Intuitively, Assumption 3.1 defines the similarity between two tasks via a Lipschitz continuity in the context space.
Revised: Intuitively, Assumption 3.1 defines the similarity between two tasks via a Lipschitz continuity in the context space, similar Lipschitz assumption also appears in Abbasi-Yadkori and Neu (2014); Modi et al. ; Dann et al.

Original:  Assumption 3.1 and 3.2 together quantify the maximum difference between two consecutive tasks M ω k − 1 , M ω k , in terms of the maximum difference between their reward functions.
Revised: ; Modi et al. Belogolovsky et al. Assumption 3. and 3.2 together quantify the maximum difference between two consecutive tasks M ω k − 1 , M ω k , in terms of the maximum difference between their reward function.

Revised: Assumption 3.1 and 3.2 also play a crucial role in reducing the exponential complexity to a polynomial one, and we will briefly discuss the intuition in the next section.

Original: 3. P RIOR A NALYSIS ON PG WITH STOCHASTIC GRADIENT
Revised: 3. P RIOR A NALYSIS ON PG WITH STOCHASTIC GRADIENT

Original: Ding et al.
Revised: Ding et al.

Original: (2021) proposed a two-phased PG convergence analysis framework with  stochastic gradient.
Revised: (2021) proposed a two-phased PG convergence analysis framework with a stochastic gradient.

Original: In particular, the author demonstrate that with high probability, stochastic PG with arbitrary initialization achieves an ε -optimal policy can be achieved with an iteration number of T 1 , T 2 and per iteration sample complexity of B 1 , B 2 in two separate phases where T 1 = (cid:101) Ω( S 2 S 3 ) ,T 2 =
Revised: In particular, the author demonstrates that with high probability, stochastic PG with arbitrary initialization achieves an ε -optimal policy can be achieved with an iteration number of T 1 , T 2 and per iteration sample complexity of B 1 , B 2 in two separate phases where T 1 = (cid:101) Ω( S 2 S 3 ) ,

Original:  (cid:101) Ω ( S 3 / 2 )and B 1 =(
Revised: T 2 = (cid:101) Ω( S 3 / 2 ) ( (cid:101) Ω( · ) suppresses the log S and terms that do not contain S ) and B 1 =

Original: cid:101) Ω( S 2 S 3 ) , B 2 = (cid:101) Ω( S 5 ) , respectively, and PG enters phase 2 only when the updating policy becomes ε 0 -optimal, where ε 0 is a term depending on S (formally defined by (19) in Appendix A.3).
Revised: (cid:101) Ω( S 2 S 3 ) , B 2 = (cid:101) Ω( S 5 ) , respectively, and PG enters phase 2 only when the updating policy becomes ε 0 -optimal, where ε 0 is a term depending on S (formally defined by (19) in Appendix A.3).

Original: For completeness, we restate the main theorem of Ding et al.
Revised: For completeness, we restate the main theorem of Ding et al.

Original: (2021) in Theorem A.2, provide the details of such dependencies on S are provided in Corollary A.3, and describes the two phase procedure in Algorithm 4.
Revised: (2021) in Theorem A.2, provide the details of such dependencies on S are provided in Corollary A.3, and describes the two phase procedure in Algorithm 4.

Original: The main implication of this two phase results is that, if we want to apply PG to learn an optimal policy from an arbitrary initialization, we suffer from an exponential iteration number and sample complexity, unless the initialization is ε 0 -optimal.
Revised: The main implication of these two-phase results is that, if we want to apply PG to learn an optimal policy from an arbitrary initialization, we suffer from an exponential iteration number and sample complexity, unless the initialization is ε 0 -optimal.

Original: In the next section, we will discuss howa curriculum satisfying Assumption 3.1 and 3.2 enables an ε 0 -optimal initialization for every ω k , and provides an overall complexity for learning the target context ω K .
Revised: In the next section, we will discuss how Assumption 3.1 and 3.2 enables an ε 0 -optimal initialization for every ω k hence reducing the exponential complexity to a polynomial one.

Original:  T HEORETICAL A NALYSIS
Revised:  T HEORETICAL A NALYSIS

Original: In this section, we introduce R OLLIN , a simple algorithm that accelerates policy learning under the contextual MDP setup by bootstrapping new context learning with a better initial distribution (Algorithm 1).
Revised: In this section, we introduce R OLLIN , a simple algorithm that accelerates policy learning under the contextual MDP setup by bootstrapping new context learning with a better initial distribution (Algorithm 1).

Original: We also provide the total complexity analysis of applying R OLLIN to stochastic PG for achieving an ε -optimal policy.
Revised: We also provide the total complexity analysis of applying R OLLIN to stochastic PG for achieving an ε -optimal policy.

Original: 4. R OLLIN
Revised: 4. R OLLIN

Original: The theoretical version of R OLLIN is provided in Algorithm 1.
Revised: The theoretical version of R OLLIN is provided in Algorithm 1.

Original: Our intuition behind R OLLIN is that when two consecutive contexts in the curriculum { ω k } Kk =0 are close, their optimal parameters θ ⋆ω k − 1 , θ ⋆ω k should be close to each other.
Revised: Our intuition behind R OLLIN is that when two consecutive contexts in the curriculum { ω k } Kk =1 are close, their optimal parameters θ ⋆ω k − 1 , θ ⋆ω k should be close to each other.

Original: Let θ ( k ) t denote the parameters at the t th iteration of stochastic PG for learning θ ⋆ω k , if we initialize the parameter θ ( k )as the optimal parameter of the previous context θ ⋆ω k − 1 (line 4 in Algorithm 1), and set the initial distribution µ k as a mixture between the optimal state visitation distribution of the previous context d π ⋆ωk − 1 µ k − 1 and the original distribution of interest ρ with ratio β (line 5 in Algorithm 1), such that
Revised: Let θ ( k ) t denote the parameters at the t th iteration of stochastic PG for learning θ ⋆ω k , if we initialize the parameter θ ( k )as the optimal parameter of the previous context θ ⋆ω k − 1 (line 5 in Algorithm 1), and set the initial distribution µ k as a mixture between the optimal state visitation distribution of the previous context d π ⋆ωk − 1 µ k − 1 and the original distribution of interest ρ with ratio β (line 6 in Algorithm 1), such that

Original: [Equation]
Revised: [Equation]

Original:  then stochastic PG enjoys  faster convergence rate.
Revised: ) then stochastic PG enjoys a faster convergence rate.

Original: This happens because setting θ ( k )
Revised: This happens because setting θ ( k )

Original: =
Revised: =

Original: θ ⋆k − 1 ensures a near-optimal initialization for θ ( k ) 1 , and setting µ k by mixing d π ⋆ωk − 1 µ k − 1 and ρ further improves the rate of convergence by decreasing the density mismatch ratio (cid:13)(cid:13)(cid:13)
Revised: θ ⋆k [Equation] ensures a near-optimal initialization for θ ( k ) 1 , and setting µ k by mixing d π ⋆ωk − 1 µ k − 1 and ρ further improves the rate of convergence by decreasing the density mismatch ratio (cid:13)(cid:13)(cid:13)

Original: d π ⋆ωk µ k /µ k (cid:13)(cid:13)(cid:13) ∞
Revised: d π ⋆ωk µ k /µ k (cid:13)(cid:13)(cid:13) ∞

Original: (a term with an S dependency that influences the rate of convergence).
Revised: (a term with an S dependency that influences the rate of convergence).

Original: [-1: Input: ρ , { ω k } Kk =0 , M W , β ∈ (0 , 1) , θ ⋆ω 0 . 2: Initialize µ 
-]
Revised: [-1: Input: ρ , { ω k } Kk =0 , M W , β ∈ (0 , 1) , θ ⋆ω 0 . 2: Initialize µ 
-]

Original: = ρ , set θ (1)1 = θ ⋆ω 0 . 3: for k = 1 , . . . , K do 4: Set θ ( k )= θ ⋆ω k − 1 .
Revised: [Equation]

Original: ▷ π ⋆ω k − 1 = 
Revised: ▷ π [Equation] = π

Original: π θ ⋆ωk− 1 is the optimal policy of ω k − 1 .
Revised: [Equation] ⋆ωk

Original: Set µ k =
Revised: [Equation]

Original: βd
Revised: 1 is

Original: π ⋆ωk − 1 µ k − 1 + (1 − β ) ρ .
Revised: [Equation] [Equation] − β ) ρ .

Original: ▷ µ k is the initial distribution for learning θ ⋆ω k .
Revised: ▷ µ k is the initial distribution for learning θ ⋆ω k .

Original: Use stochastic PG (Algorithm 4) with initialization θ ( k ) 1 , µ k , M ω k and obtain θ ⋆ω k .
Revised: Run stochastic PG (Algorithm 4) with initialization θ ( k ) 1 , µ k , M ω k and obtain θ ⋆ω k .

Original: 7: end for 8: Output: θ ⋆ω K
Revised: 8: end for 9: Output: θ ⋆ω K

Original: 4. M AIN R ESULTS
Revised: 4. M AIN R ESULTS

Original: We now discuss how to use a sequence of contexts to learn the target context ω K with provable efficiency given the optimal policy π ⋆ω 0 of the initial context ω 0 , without incurring an exponential dependency on S (as mentioned in Section 3.2).
Revised: We now discuss how to use a sequence of contexts to learn the target context ω K with provable efficiency given the optimal policy π ⋆ω 0 of the initial context ω 0 , without incurring an exponential dependency on S (as mentioned in Section 3.2).

Original: Our polynomial complexity comes as a result of enforcing an ε 0 -optimal initialization ( ε 0 is the same as Section 3.2 and (19)) for running stochastic
Revised: Our polynomial complexity comes as a result of enforcing an ε 0 -optimal initialization ( ε 0 is the same as Section 3.2 and (19)) for running stochastic

Original: PG (line 6 of Algorithm 1).
Revised: PG (line 6 of Algorithm 1).

Original: Hence, stochastic PG directly enters phase 2, with a polynomial dependency in S .
Revised: Hence, stochastic PG directly enters phase 2, with a polynomial dependency in S .

Original: Our main results consists of two parts.
Revised: Our main results consist of two parts.

Original: We first show that when two consecutive contexts ω k − 1 , ω k are close enough to each other, using R OLLIN for learning θ ⋆k with initialization θ ( k )
Revised: We first show that when two consecutive contexts ω k − 1 , ω k are close enough to each other, using R OLLIN for learning θ ⋆k with initialization θ ( k )

Original: = θ ⋆ω k − 1 and applying an initial distribution µ k = βd
Revised: = θ ⋆ω k − 1 and applying an initial distribution µ k = βd

Original: π ⋆ωk − 1 µ k − 1 +(1 − β ) ρ improves the convergence rate.
Revised: π ⋆ ω k − 1 µ k − 1 +(1 − β ) ρ improves the convergence rate.

Original: Specifically, the iteration number and complexity for learning θ ⋆ω
Revised: Specifically, the iteration number and complexity for learning θ ⋆ω

Original: k from θ ⋆ω k − 1 is stated as follows:
Revised: k from θ ⋆ω k − 1 is stated as follows:

Original: Theorem 4.1 (Complexity of Learning the Next Context) Consider the context based stochastic softmax policy gradient (line 6 of Algorithm 1), if ∥ ω k − ω k − 1 ∥ 2 ≤ O (cid:0) S − 2 (cid:1) , then the iteration number of obtaining an ε -optimal policy for ω k from θ ⋆ω k − 1 is (cid:101) Ω (cid:0) S 3 / 2 (cid:1) and the per iteration sample complexity is (cid:101) Ω (cid:16) L r α (1 − β ) S 3 (cid:17) .
Revised: Theorem 4.1 (Complexity of Learning the Next Context) Consider the context-based stochastic softmax policy gradient (line 7 of Algorithm 1), suppose Assumption 3.1 and Assumption 3.2 hold, then the iteration number of obtaining an ε -optimal policy for ω k from θ ⋆ω k − 1 is (cid:101) Ω (cid:0) S 3 / 2 (cid:1) and the per iteration sample complexity is (cid:101) Ω (cid:16) L r α (1 − β ) S 3 (cid:17) .

Original: In other words, Theorem 4.1 shows that when ω k − 1 , ω k are close enough, R OLLIN reduces the iteration and sample complexity from an exponential dependency of (cid:101) Ω( S 2 S 3 ) to an iteration number of (cid:101) Ω( S 3 / 2 ) and per iteration sample complexity of (cid:101) Ω( S 3 ) .
Revised: In other words, Theorem 4.1 shows that when ω k − 1 , ω k are close enough, R OLLIN reduces the iteration and sample complexity from an exponential dependency of (cid:101) Ω( S 2 S 3 ) to an iteration number of (cid:101) Ω( S 3 / 2 ) and per iteration sample complexity of (cid:101) Ω( S 3 ) .

Original: It is worth noting that the theorem above only address the iteration numbers and sample complexity for learning θ ⋆ω k from θ ⋆ω k − 1 .
Revised: It is worth noting that the theorem above only address the iteration numbers and sample complexity for learning θ ⋆ ω k from θ ⋆ ω k − 1 .

Original: The next theorem provides the total complexity for learning θ ⋆ω
Revised: Theorem 4.3 provides the total complexity for learning θ ⋆ω

Original: K from θ ⋆ω 0 via recursively apply the results in 
Revised: K from θ (0) via recursively apply the results in Theorem 4.1.

Original:  Theorem 4.1.
Revised: Before introducing Theorem 4.3, we first provide a criteria for the desired initialization of θ (0)0 . Definition 4.2 (Near-optimal Initialization) We say θ 0 is a near-optimal initialization for learning θ ⋆ω

Original: Concretely, we have:
Revised: if θ 0 satisfies [Equation] Note that in the above definition, π ⋆ω k represents the optimal policy of ω k , and V πω k represents value function of context ω k under policy π . Now we introduce the results for the overall complexity:

Original: Theorem 4.2 (Main Results: Total Complexity of R OLLIN )
Revised: Theorem 4.3 (Main Results: Total Complexity of R OLLIN )

Original: Suppose the assumptions in Theorem 4.1 hold.
Revised: Suppose Assumption 3.1 and 3.

Original: Given a curriculum { ω k } Kk =0 such that max 1 ≤ k ≤ K ∥ ω k − ω k − 1 ∥ 2 ≤ O ( S − 2 ) , the total number of iteration of learning π ⋆ω K using Algorithm 1 is Ω( KS 3 / 2 ) and the per iteration is (cid:101) Ω (cid:0) S 3 (cid:1) , with high probability.
Revised: hold, and θ (0) is an near-optimal initialization, then the total number of iteration of learning π ⋆ω K using Algorithm 1 is Ω( KS 3 / 2 ) and the per iteration is (cid:101) Ω (cid:0) S 3 (cid:1) , with high probability.

Original: A direct implication of Theorem 4.2 is that, with a curriculum { ω k } Kk =0 satisfying Assumption 3.1 and 3.2, one can reduce the daunting exponential dependency on S caused by poor initialization to a polynomial dependency on S .
Revised: A direct implication of Theorem 4.3 is that, with a curriculum { ω k } Kk =0 satisfying Assumption 3.1 and 3.2, one can reduce the daunting exponential dependency on S caused by poor initialization to a polynomial dependency on S .

Original: Admittedly the state space S itself is still large in practice, but reducing the state space S itself requires extra assumptions on S which is beyond the scope of this work.
Revised: Admittedly the state space S itself is still large in practice, but reducing the state space S itself requires extra assumptions on S which is beyond the scope of this work.

Original: We will provide a sketch proof of Theorem 4.1 and 4.2 in the next subsection and leave all the details proof to Appendix A.4, A.5 respectively.
Revised: We will provide a sketch proof of Theorem 4.1 and 4.3 in the next subsection and leave all the details proof to Appendix A.4, A.5 respectively.

Original: 4. P ROOF S KETCHSketch
Revised: 4. P ROOF S KETCH

Original:  proof of Theorem 4.1 The key insight for proving Theorem 4.1 is to show thatR OLLIN provides a good initialization is to show that in MDP M ω k , the value function w.r.t.
Revised: Sketch proof of Theorem 4.1 The key insight for proving Theorem 4.1 is to show that in MDP M ω k , the value function w.r.t.

Original: π ⋆ ω k , π ⋆ ω k −can be bounded by the ℓ 2 norm between ω k and ω k − 1 .
Revised: π ⋆ω k , π ⋆ω k − 1 can be bounded by the ℓ 2 norm between ω k and ω k − 1 .

Original: In particular, we prove such a relation in Lemma A.5:
Revised: In particular, we prove such a relation in Lemma A.5:

Original: [Equation]
Revised: [Equation]

Revised: By setting θ ( k )

Original: By setting θ ( k )= θ ⋆ω k − 1 , Equation (2) directly implies Vπ ⋆ωk
Revised: [Equation] θ [Equation] 1 , Equation (2) directly implies V

Original: [Equation]
Revised: [Equation]

Original: As suggested by Ding et al.
Revised: As suggested by Ding et al.

Original: (2021) stochastic PG can directly start from stage 2 with polynomial complexity of T 2 = 
Revised: (2021) stochastic PG can directly start from stage 2 with polynomial complexity of T [Equation] = [Equation]

Revised: Ω

Revised: [Equation]

Original: (cid:101) Ω( S 3 / 2 ) , B[-2 = (cid:101) Ω( S 5 ) , if V π ⋆ωk ω k ( ρ ) − V θ ( k ) 
-]
Revised: S [Equation] [Equation] ) , B[-2 = (cid:101) Ω( S 5 ) , if V π ⋆ωk ω k ( ρ ) − V θ ( k ) 
-]

Revised: [Equation]

Revised: =

Revised: [Equation]

Revised: Ω

Revised: [Equation]

Revised: S 5 ) , if V π ⋆ωk ω k

Original:  ( ρ ) ≤ ε 0 , where ε 0 (formally defined in Equation (19) in Appendix A.3) is a constant satisfying ε 

Revised: [Equation] ρ ) − [Equation] k ) 1 ( ρ ) ≤ ε 0 , where ε 0 (formally defined in Equation (19) in Appendix A.3) is a constant satisfying ε

Original: =
Revised: =

Original: O ( S − 2 ) .
Revised: O ( S − 2 ) .

Original: Hence, by enforcing two consecutive contexts are close enough ∥ ω k − ω k − 1 ∥ 2 ≤ O ( S − 2 ) , we can directly start from a near-optimal initialization with polynomial complexity w.r.t.
Revised: Hence, by enforcing two consecutive contexts to be close enough ∥ ω k − ω k − 1 ∥ 2 ≤ O ( S − 2 ) , we can directly start from a near-optimal initialization with polynomial complexity w.r.t.

Original: S .
Revised: S .

Original: This largely explains the intuition behind the initialization in line 4 of R OLLIN :
Revised: This largely explains the intuition behind the initialization in line 5 of R OLLIN :

Original: θ ( k )=
Revised: θ ( k )=

Original: θ ⋆ω k − 1 .
Revised: θ ⋆ω k − 1 .

Original: It is worth highlighting that the per iteration sample complexity B 2 shown by Ding et al.
Revised: It is worth highlighting that the per iteration sample complexity B 2 shown by Ding et al.

Original: (2021) scales as (cid:101) Ω( S 5 ) , while our result in
Revised: (2021) scales as (cid:101) Ω( S 5 ) , while our result in

Original: Theorem 4.1 only requires a smaller sample complexity of (cid:101) Ω( S 3 ) .
Revised: Theorem 4.1 only requires a smaller sample complexity of (cid:101) Ω( S 3 ) .

Original: Such an improvement in the sample complexity comes from line 5 of R OLLIN :
Revised: Such an improvement in the sample complexity comes from line 6 of R OLLIN :

Original: µ k = βd
Revised: µ k = βd

Original: π ⋆ωk − 1 µ k − 1 + (1 − β ) ρ .
Revised: π ⋆ωk − 1 µ k − 1 + (1 − β ) ρ .

Original: Intuitively,by setting µ k as βd π ⋆ωk − 1 µ k − 1 + (1 − β ) ρ allows us to provide an upper bound on the density mismatch ratio:
Revised: Intuitively, setting µ k as βd π ⋆ωk − 1 µ k − 1 + (1 − β ) ρ allows us to provide an upper bound on the density mismatch ratio:

Original: [Equation]
Revised: [Equation]

Original: where ∆ kω = max 1 ≤ i ≤ k ∥ ω
Revised: where ∆ kω = max 1 ≤ i ≤ k ∥ ω

Original: i − ω
Revised: i − ω

Original: i − 1 ∥ 2 .
Revised: i − 1 ∥ 2 .

Original: Since the sample complexity B 2 (provided in Corollary A.3) contains one multiplied of (cid:13)(cid:13) (cid:13)(cid:13)
Revised: Since the sample complexity B 2 (provided in Corollary A.3) contains one multiplier of (cid:13)(cid:13)(cid:13)(cid:13)

Original: [Equation] ∞ , by assuming ∆ kω
Revised: d π ⋆µk µ k /µ k (cid:13)(cid:13)(cid:13)(cid:13) ∞ , setting ∆ kω

Original: = O ( S − 2 ) immediately reduces the complexity by an order of S 2 .
Revised: = O ( S − 2 ) immediately reduces the complexity by an order of S 2 .

Original: The proof of the upper bound of the density mismatch ratio (Equation 3) is provided in Lemma A.1.
Revised: The proof of the upper bound of the density mismatch ratio (Equation 3) is provided in Lemma A.1.

Original: Sketch proof of Theorem 4.2 We obtain Theorem 4.2 by recursively applying Theorem 4.1.
Revised: Sketch proof of Theorem 4.3 We obtain Theorem 4.3 by recursively applying Theorem 4.1.

Original: More precisely, we usean induction to show
Revised: More precisely, we use induction to show

Original: that, if we initialize the parameters of the policy as θ ( k )= θ ⋆ω k − 1 , when t = (cid:101) Ω( S 3 / 2 ) ,
Revised: that, if we initialize the parameters of the policy as θ ( k )= θ ⋆ω k − 1 , when t = (cid:101) Ω( S 3 / 2 ) ,

Original: ∀ k ∈
Revised: ∀ k ∈

Original: [ K ] , we have V π ⋆ωk ω k ( ρ ) − V π [Equation]
Revised: [ K ] , we have V π ⋆ωk ω k ( ρ ) − V π θ ( k − 1) t ω k

Original: ( ρ ) < ε 0 .
Revised: ( ρ ) < ε 0 .

Original: This implies that for any context ω k , k ∈ [ K ] , initializing θ ( k )
Revised: Hence, for any context ω k , k ∈ [ K ] , initializing θ ( k )

Original: [Equation]
Revised: =

Original:  from learningthe previous context π ⋆ω [Equation] using stochastic PG after t=
Revised: θ ( k − 1) t from learning π ⋆ω k − 1 via stochastic PG after t

Original:  Ω( S 3 / 2 ) iteration, θ ( k )will directly start from the efficient phasewith a polynomial complexity.
Revised: = Ω( S 3 / 2 ) iteration, θ ( k ) 1 will directly start from the efficient phase 2 with polynomial complexity.

Original: Hence, the total iteration number for learning the θ ⋆ K is Ω( KS 3 / 2 ) , and the per iteration sample complexity (cid:101) Ω (cid:0) S 3 (cid:1) remains the same as Theorem 4.1.
Revised: Hence, the total iteration number for learning the θ ⋆K is Ω( KS 3 / 2 ) , and the per iteration sample complexity (cid:101) Ω (cid:0) S 3 (cid:1) remains the same as Theorem 4.1.

Original:  P RACTICAL I MPLEMENTATION OF R OLLIN
Revised:  P RACTICAL I MPLEMENTATION OF R OLLIN

Original: In this section, we describe the practical implementation of R OLLIN using Soft-Actor-Critic (Haarnoja et al., 2018)
Revised: In this section, we describe the practical implementation of R OLLIN using Soft-Actor-Critic (Haarnoja et al., 2018)

Original: algorithm when an oracle curriculum is available.
Revised: algorithm when an oracle curriculum is available.

Original: SAC can be seen as a variant of entropy-regularized stochastic PG with the addition of the critics to reduce the gradient variance.
Revised: SAC can be seen as a variant of entropy-regularized stochastic PG with the addition of the critics to reduce the gradient variance.

Original: Recall that in the theoretical analysis, we learn a separate policy for each context that can start from the near-optimal state distribution of the previous context to achieve  good return under the current context.
Revised: Recall that in the theoretical analysis, we learn a separate policy for each context that can start from the near-optimal state distribution of the previous context to achieve a good return under the current context.

Original: However, in practice, we usually would want to have a policy that can directly start from the initial distribution ρ to obtain  good return for the final context ω K .
Revised: However, in practice, we usually would want to have a policy that can directly start from the initial distribution ρ to obtain a good return for the final context ω K .

Original: In order to learn such  policy, we propose to have two context-conditioned RL agents training in parallel, where the first agent π main is the main agent that eventually will learn to achieve a good return from ρ , and the second agent π exp is an exploration agent that learns to achieve a good return under the current context from the near-optimal state density of the previous context.
Revised: In order to learn such a policy, we propose to have two context-conditioned RL agents training in parallel, where the first agent π main is the main agent that eventually will learn to achieve a good return from ρ , and the second agent π exp is an exploration agent that learns to achieve a good return under the current context from the near-optimal state density of the previous context.

Original: Another purpose of the exploration agent (as the name suggests) is to provide  better exploration experience for the main agent to learn the current context better.
Revised: Another purpose of the exploration agent (as the name suggests) is to provide a better exploration experience for the main agent to learn the current context better.

Original: This is made convenient by using  off-policy RL agent where the main agent can easily learn from data that is not on-policy.
Revised: This is made convenient by using an off-policy RL agent where the main agent can easily learn from data that is not on-policy.

Original: Specifically, for each episode, there is a probability of β where we run the main agent conditioned on the previous context for a random number of steps until we switch to the exploration agent to collect experience for the current context until the episode ends.
Revised: Specifically, for each episode, there is a probability of β where we run the main agent conditioned on the previous context for a random number of steps until we switch to the exploration agent to collect experience for the current context until the episode ends.

Original: Otherwise, we directly run the main agent for the entire episode.
Revised: Otherwise, we directly run the main agent for the entire episode.

Original: Both agents are trained to maximize the return under the current context.
Revised: Both agents are trained to maximize the return under the current context.

Original: Whenever the average return of the last 10 episodes exceed a performance threshold R , we immediately switch to the next context and re-initialize the exploration agent and its replay buffer.
Revised: Whenever the average return of the last 10 episodes exceed a performance threshold R , we immediately switch to the next context and re-initialize the exploration agent and its replay buffer.

Original: A high-level description is available in Algorithm 2 (a more detailed version in Algorithm 8).
Revised: A high-level description is available in Algorithm 2 (a more detailed version in Algorithm 8).

Original: Algorithm 0 and with probability of β then 8: h ∼ Geom(1 − γ ) (truncated at H ) 9: run π main ( a | s, ω k − 1 ) from the initial state for h steps and switch to π exp ( a | s, ω k ) until the episode ends to obtain trajectory τ 0:
Revised: s, ω k − 1 ) from the initial state for h steps and switch to π exp ( a | s, ω k ) until the episode ends to obtain trajectory τ 0:

Original: H = { s 0 , a 0 , r 0 , s 1 , a 1 , · · · , s H } .
Revised: H = { s 0 , a 0 , r 0 , s 1 , a 1 , · · · , s H } .

Original: 10: record τ 0: H in D , and τ h : H in D exp .
Revised: 10: record τ 0: H in D , and τ h : H in D exp .

Original: 11: else 12: run π main ( a | s, ω k ) to obtain trajectory τ 0: H and record τ 0: H in D . 13: end if 14: at each environment step during the episode, update π main using D and π exp using D exp .
Revised: 11: else 12: run π main ( a | s, ω k ) to obtain trajectory τ 0: H and record τ 0: H in D . 13: end if 14: at each environment step during the episode, update π main using D and π exp using D exp .

Original: 15: end for 16: Output: π main
Revised: 15: end for 16: Output: π main

Original:  E XPERIMENTAL R ESULTS
Revised:  E XPERIMENTAL R ESULTS

Original: κ = 2/ κ = 1/
Revised: κ = 2/ κ = 1/

Original: While the focus of our work is on developing a provably efficient approach to curriculum learning, we also conduct an experimental evaluation of our practical implementation of R OLLIN with soft actorcritic (SAC) (Haarnoja et al., 2018) as the RL algorithm on  the antmaze-umaze environment (Fu et al., 2020).
Revised: While the focus of our work is on developing a provably efficient approach to curriculum learning, we also conduct an experimental evaluation of our practical implementation of R OLLIN with soft actor-critic (SAC) (Haarnoja et al., 2018) as the RL algorithm on several continuous control tasks including goal reaching with an oracle curriculum and a learned curriculum, as well as non-goal tasks. 6. G OAL R EACHING WITH AN O RACLE C URRICULUM We adopt the antmaze-umaze environment (Fu et al., 2020) for evaluating the performance of R OLLIN in goal-reaching tasks.

Original: We use a hand-crafted path of contexts, where each context specifies the location that the ant needs to reach (as shown in Figure 2).
Revised: In the oracle curriculum case, we use a hand-crafted path of contexts, where each context specifies the location that the ant needs to reach (as shown in Figure 2).

Original: 2 We consider a path of contexts ω
Revised: 1 We consider a path of contexts ω

Original: ( κ ) parameterized by κ ∈ [0 , 1] where ω (0) =
Revised: ( κ ) parameterized by κ ∈ [0 , 1] where ω (0) =

Original: ω 0 and ω (1) =
Revised: ω 0 and ω (1) =

Original: ω K , and propose contexts along the path with a fixed step size ∆ .
Revised: ω K , and propose contexts along the path with a fixed step size ∆ .

Original: SeeAppendix
Revised: See

Original: E for more implementation details.
Revised: Appendix E.1 for more implementation details.

Original: Wecan combine R OLLIN with a variety of prior methods, and we evaluate the following conditions: (1) standard goal reaching; (2) goal reaching with goal relabeling (Andrychowicz et al., 2017);(3) goal reaching with
Revised: We combine R OLLIN with a variety of prior methods, and we evaluate the following conditions: (1) standard goal reaching; (2) goal reaching with goal relabeling (Andrychowicz et al., 2017);

Original:  Go-Explore (Ecoffet et al., 2019).
Revised: (3) goal reaching with Go-Explore (Ecoffet et al., 2019).

Original: For goal relabeling , we adopt a similar relabeling technique as Pitis et al.
Revised: For goal relabeling, we adopt a similar relabeling technique as Pitis et al.

Original: (2020), where each minibatch contains 1 / 3 original transitions 1 / 3 transitions with future state relabeling , and 1 / 3 transitions with next state relabeling.
Revised: (2020), where each mini batch contains 1 / 3 original transitions, 1 /transitions with future state relabeling, and 1 / 3 transitions with next state relabeling.

Original: We implemented the Go-Explore method by adding an additional standard Gaussian exploration noise (multiplied by a constant factor) to the agent for learning the next goal ω ( k + 1) , once it reaches the current goal ω ( k ) .
Revised: We implemented the Go-Explore method by adding an additional standard Gaussian exploration noise (multiplied by a constant factor) to the agent for learning the next goal ω ( k + 1) , once it reaches the current goal ω ( k ) .

Original: We empirically observed that sampling the replay buffer from a geometric distribution with p = 10 − 5 (more recent transitions are sampled more frequently) improves the overall performance.
Revised: We empirically observed that sampling the replay buffer from a geometric distribution with p = 10 − 5 (more recent transitions are sampled more frequently) improves the overall performance.

Original: Hence, in all future experiments, we compare  performance of R OLLIN with classic uniform sampling and the new geometric sampling.
Revised: Hence, in all future experiments, we compare the performance of R OLLIN with classic uniform sampling and the new geometric sampling.

Original: We demonstrate how adjusting the R OLLIN parameter ( β = 0 . 1 , 0 . 2 , 0 . 5 , 0 .
Revised: We demonstrate how adjusting the R OLLIN parameter ( β = 0 . 1 , 0 . 2 , 0 . 5 , 0 .

Original: 75 , 0 . 9 ) impacts the learning speed on three different step sizes ∆ = 124 , 118 , 112 .
Revised: 75 , 0 . 9 ) impacts the learning speed on three different step sizes ∆ = 124 , 118 , 112 .

Original: We expect coarser curricula with larger steps to be more difficult.
Revised: We expect coarser curricula with larger steps to be more difficult.

Original: Main comparison.
Revised: Main comparisons.

Original: We first provide an overview experiments that compares R OLLIN with a fixed β = 0 .
Revised: We first provide an overview experiments that compares R OLLIN with a fixed β = 0 .

Original: 1 on different step sizes ∆ in different settings.
Revised: 1 on different step sizes ∆ in different settings.

Original: In each case we compare the prior method (vanilla, relabeled, or Go-Explore) with and without the addition of R OLLIN .
Revised: In each case, we compare the prior method (vanilla, relabeled, or Go-Explore) with and without the addition of R OLLIN .

Original: Note that for the case without R OLLIN , we still provide the curriculum to the agent for training.

Original: As shown in Table 1,κ progresses further with R OLLIN  in most presented settings (except Go-Explore with ∆ = 1 / 12 ).
Revised: As shown in Table 1, R OLLIN improves the largest value of κ reached by the agent in most presented settings (except Go-Explore with ∆ = 1 / 12 ).

Original: This result suggests that R OLLIN facilitates goal-conditioned RL with a curriculum, asthe we only update the learning progress κ to κ +∆ when the return of the current policy reaches  certain threshold R (See detailed update of κ in Algorithm 2).
Revised: This result suggests that R OLLIN facilitates goal-conditioned RL with a curriculum, as we only update the learning progress κ to κ + ∆ when the return of the current policy reaches a certain threshold R (See detailed update of κ in Algorithm 2).

Original: Since geometric sampling technique drastically improvement in the vanilla and the Go-Explore settings and slightly decreases the performance in the relabeling setting, we also want to point out that R OLLIN improves the learning progress of baselines at different levels (w/ and w/o geometric samplings).

Original: Note that β = 0 .
Revised: Note that β = 0 .

Original: 1 does not always produce the best result, we will provide more results comparing different β s in different settings later in this section
Revised: 1 does not always produce the best result, we will provide more results comparing different β s in different settings later in this section, and we leave all the learning curves and detailed tables to Appendix F.1.

Revised: Note that we do not include the results of directly learning the last context in the antmaze-umaze environment because the agent cannot reach the goal without the aid of a curriculum, which is corroborated by Pitis et al.

Original: [Table]
Revised: [Table]

Original: and we leave all the learning curves to Appendix F
Revised: 6. G OAL R EACHING WITH A L EARNED C URRICULUM

Original: [Table]
Revised: [Figure]

Revised: In this section, we focus on the setting where an oracle curriculum is not provided.

Revised: In particular, we show that R OLLIN can improve MEGA (Pitis et al., 2020), an existing automated goal curriculum generation method (Figure 3).

Revised: MEGA proposes intrinsic goals that are achievable, but appear rarely in the replay buffer.

Revised: Since there is no clear notion of progression in the learned curriculum setting, we use the progression in the number of environment steps where ω k corresponds to the goals during k × 100 K - ( k + 1) × 100 K steps.

Revised: We follow the same highlevel procedure as described in Algorithm 2 where we randomly roll in a policy that was trained on the previous context (in this case, it would be the policy checkpoint that was obtained after the last 100K-step learning segment).

Original: Goal reaching in different settings.
Revised: A detailed algorithm procedure is provided in Algorithm 9. 6. N ON -G OAL R EACHING T ASKS For the non-goal tasks, we choose a fixed contextual space with ten discrete contexts: κ ∈ { 0 . 1 , 0 . 2 , . . , 1 } , where each κ uniquely determines the desired x -velocity of a locomotion agent in the following environments: walker2d , hopper , humanoid , and ant in OpenAI gym

Revised: (Brockman et al., 2016).

Revised: Our goal is to train an agent to move fast and stably.

Original: For the vanilla goal reaching case, we present the final κ values of 5 β s after 3 million environment steps in Table 2.
Revised: For each context ω ( κ ) , we set the desired speed range to be [ λκ, λ ( κ + 0 . 1)) , where λ is a parameter depending on the physics of the agent in different environments. When the x -velocity of an agent is within the desired velocity range, we set the healthy_reward to a higher value healthy_reward_high . Otherwise, we set the healthy_reward to a lower value healthy_reward_low .

Original: Across all step sizes, R OLLIN is able to progress further in κ consistently compared to the baseline, demonstrating the effectiveness of
Revised: In each environment, we increase the task difficulty with later curriculum steps (larger κ ), by increasing the near-optimal threshold R ( κ ) .

Original:  R OLLIN in this curriculum learning setup. In particular, the first three rows of Table 2 show that
Revised: Detailed parameters of the desired speed range λ , near optimal-threshold R ( κ ) , healthy_reward_high , and healthy_reward_low are provided in Appendix E.3.

Original:  R OLLIN achieves much better results than the baseline method on all tested β s, without geometric sampling.
Revised: Main comparisons. We first compare R OLLIN with a fixed β = 0 . 1 at different environment steps: 0 . 75 × 10 6 , 1 × 10 6 .

Original: Note that the bottom three rows of Table 2 also suggests that  R OLLIN largely improves the results on a better the baseline with geometric sampling (except for some extreme β s). We also observe similar phenomenon in  the relabeling and Go-Explore settings (see details in Table 3, 4)
Revised: In each case, we compare the learning progress κ , averaged x -velocity, and averaged return, with and without the addition of R OLLIN . Note that for the case without R OLLIN , we still provide the curriculum to the agent for training. As shown in Table 2, R OLLIN improves the largest learning progresses κ , average x -velocity, and average return in most presented settings.

Original: [Table]
Revised: This result suggests that R OLLIN also can facilitate learning of non-goal tasks, as we only update the learning progress κ to κ + 0 .

Original: [Table]
Revised: 1 when the return of the current policy reaches a certain threshold R ( κ ) (See detailed update of κ in Algorithm 2). Note that β = 0 . 1 does not always produce the best result, we will provide more results comparing different β s in different settings later in this section, and we leave all the learning curves and detailed tables to Appendix F.3.

Original: [Table]
Revised: [Figure]

Original: In summary

Original: [Table]
Revised: [Table]

Original:  we empirically demonstrated that R OLLIN improves the performance of goal reaching tasks in different settings with a vast range of selection of β s. Although R OLLIN introduces an extra parameter β for finetuning goal reaching tasks with a curriculum, our extensive experiments suggest that one could expect improvement by choosing a constant β = 0 .
Revised: In summary, we empirically demonstrated that R OLLIN improves the performance of goal-reaching tasks and non-goal tasks in different settings with a vast range of selection of β s. Although R OLLIN introduces an extra parameter β for finetuning goal-reaching tasks with a curriculum, our extensive experiments suggest that one could expect improvement by choosing a constant β = 0 .

Original: 1 or 0 . [-
-]
Revised: 1 or 0 . [-
-]2 .

Original: for other goal reaching tasks. Still, for achieving the best performance in general task, one shall still consider finetuning the R OLLIN parameter β .
Revised: Still, for achieving the best performance in general-task, one shall still consider fine-tuning β .

Original:  D ISCUSSION AND L IMITATIONS
Revised:  D ISCUSSION AND L IMITATIONS

Original: We presented R OLLIN , a simple algorithm that accelerates curriculum learning under the contextual
Revised: We presented R OLLIN , a simple algorithm that accelerates curriculum learning under the contextual

Original: MDP setup by rolling in a near optimal policy to bootstrap the learning of new nearby contexts with provable learning efficiency benefits.
Revised: MDP setup by rolling in a near-optimal policy to bootstrap the learning of new nearby contexts with provable learning efficiency benefits.

Original: Our focus is on the theoretical analysis of R OLLIN , which analyzes the method instantiated with entropy-regularized policy gradients and a softmax policy parameterization.

Original: We show that R OLLIN attains polynomial sample complexity by utilizing adjacent contexts to initialize each policy.
Revised: Theoretically, we show that R OLLIN attains polynomial sample complexity by utilizing adjacent contexts to initialize each policy.

Original: Since the key theoretical insight of R OLLIN suggests that one can reduce the density mismatch ratio by constructing a new initial distribution, it would be interesting to see how R OLLIN can affect other variants of convergence analysis of PG (e.g., NPG (Kakade, 2001; 
Revised: Since the key theoretical insight of R OLLIN suggests that one can reduce the density mismatch ratio by constructing a new initial distribution, it would be interesting to see how R OLLIN can affect other variants of convergence analysis of PG (e.g., NPG (Kakade, 2001; Cen et al., 2021) or PG in a feature space (Agarwal et al., 2021; 2020)) and we would like to leave it to future. On the empirical side, our extensive experiments also suggest that one could expect improvement by choosing a constant β = 0 .

Original: Cen et al., 2021) or PG in{+0 . 
 by applying R OLLIN to various RL tasks with+} a feature space (Agarwal et al., 2021; 2020)) and we would  like to leave it to future.
Revised: 1 or {+0 . 
 by applying R OLLIN to various RL tasks with+} a curriculum. Hence, another potential direction on the empirical side would be providing an automatic tuning method for selecting the best β , which we would also like to leave for the future.

Original: A PPENDIX A G ENERALIZATION B ETWEEN D IFFERENT T ASKS IN THE C ONTEXT S PACE A. S UMMARIES OF N OTATIONS AND A SSUMPTIONS
Revised: A PPENDIX A G ENERALIZATION B ETWEEN D IFFERENT T ASKS IN THE C ONTEXT S PACE A. S UMMARIES OF N OTATIONS AND A SSUMPTIONS

Original: The maximum entropy RL (MaxEnt RL) objective with initial state distribution ρ in reinforcement aims at maximizing (Equation 15 & 16 of Mei et al
Revised: The maximum entropy RL (MaxEnt RL) objective with initial state distribution ρ in reinforcement aims at maximizing (Equation 15 & 16 of Mei et al

Original: [Equation]
Revised: [Equation]

Original: [Equation]
Revised: [Equation]

Original: and H ( π ( a h | s h )) is the discounted entropy term
Revised: and H ( π ( a h | s h )) is the discounted entropy term

Original: [Equation]
Revised: [Equation]

Original: and α is the penalty term.
Revised: and α is the penalty term.

Original: For simplicity, we denote the optimization objective function in (4) as α -MaxEnt RL.
Revised: For simplicity, we denote the optimization objective function in (4) as α -MaxEnt RL.

Original: Similar to Equation 18 & 19 of Mei et al.
Revised: Similar to Equation 18 & 19 of Mei et al.

Original: (2020), we also define the advantage and Q -functions and for MaxEnt RL as
Revised: (2020), we also define the advantage and Q -functions and for MaxEnt RL as

Original: [Equation]
Revised: [Equation]

Original: We let
Revised: We let

Original: [Equation]
Revised: [Equation]

Original: to denote the discounted state visitation of policy π starting at state s 0 , and let
Revised: to denote the discounted state visitation of policy π starting at state s 0 , and let

Original: [Equation]
Revised: [Equation]

Original: denote the initial state visitation distribution under initial state distribution ρ .
Revised: denote the initial state visitation distribution under initial state distribution ρ .

Original: We assume the reward functions under all context are bounded within [
Revised: We assume the reward functions under all context are bounded within [

Original: [Equation] [Equation]
Revised: [Equation] [Equation]

Original: Similar to previous analysis in Agarwal et al.
Revised: Similar to previous analysis in Agarwal et al.

Original: (2021); Mei et al.
Revised: (2021); Mei et al.

Original: Ding et al.
Revised: Ding et al.

Original: 2021), we assume the initial distribution µ for PG/stochastic PG satisfies ρ ( s ) > 0 , ∀ s ∈ S
Revised: 2021), we assume the initial distribution µ for PG/stochastic PG satisfies ρ ( s ) > 0 , ∀ s ∈ S

Original: A. M AIN R ESULTS : M ISMATCH C OEFFICIENT U PPER B OUND
Revised: A. M AIN R ESULTS : M ISMATCH C OEFFICIENT U PPER B OUND

Original: Lemma A.1 (Density Mismatch Ratio via R OLLIN )
Revised: Lemma A.1 (Density Mismatch Ratio via R OLLIN )

Original: Using (1) from R OLLIN µ k = βd
Revised: Using (1) from R OLLIN µ k = βd

Original: π ⋆ωk − 1 µ k − 1 + (1 − β ) ρ , the density mismatch ratio (cid:13)(cid:13)(cid:13)(cid:13)
Revised: π ⋆ωk − 1 µ k − 1 + (1 − β ) ρ , the density mismatch ratio (cid:13)(cid:13)(cid:13)(cid:13)

Original: d π ⋆ωk µ k /µ k (cid:13)(cid:13)(cid:13)(cid:13) ∞ satisfies (cid:13) (cid:13)(cid:13)(cid:13)(cid:13)(cid:13)
Revised: d π ⋆ωk µ k /µ k (cid:13)(cid:13)(cid:13)(cid:13) ∞ satisfies (cid:13) (cid:13)(cid:13)(cid:13)(cid:13)(cid:13)

Original: d π ⋆ ωk µ k µ k
Revised: d π ⋆ ωk µ k µ k

Original: (cid:13) (cid:13)(cid:13)(cid:13)(cid:13)(cid:13) ∞ ≤ (cid:101) Ω (cid:18)
Revised: (cid:13) (cid:13)(cid:13)(cid:13)(cid:13)(cid:13) ∞ ≤ (cid:101) Ω (cid:18)

Original: L r α (1 − β )∆ kω S (cid:19) , (10) where ∆ kω = max 1 ≤ i ≤ k ∥ ω
Revised: L r α (1 − β )∆ kω S (cid:19) , (10) where ∆ kω = max 1 ≤ i ≤ k ∥ ω

Original: i − ω i − 1 ∥ 2 .
Revised: i − ω i − 1 ∥ 2 .

Original: Proof By (1) from R OLLIN , we have (cid:13)(cid:13)(cid:13)(cid:13)(cid:13)(cid:13)
Revised: Proof By (1) from R OLLIN , we have (cid:13)(cid:13)(cid:13)(cid:13)(cid:13)(cid:13)

Original: d π ⋆ωk µ k µ k (cid:13)(cid:13)(cid:13)(cid:13)(cid:13)(cid:13) ∞ = (cid:13)(cid:13)(cid:13)(cid:13)(cid:13)(cid:13)
Revised: d π ⋆ωk µ k µ k (cid:13)(cid:13)(cid:13)(cid:13)(cid:13)(cid:13) ∞ = (cid:13)(cid:13)(cid:13)(cid:13)(cid:13)(cid:13)

Original: d π ⋆ωk µ
Revised: d π ⋆ωk µ

Original: [Equation]
Revised: [Equation]

Original: − d
Revised: − d

Original: [Equation]
Revised: [Equation]

Original: ⋆ωk − 1 µ k − 1 +
Revised: ⋆ωk − 1 µ k − 1 +

Original: [Equation]
Revised: [Equation]

Original: where inequality ( i ) holds because of (1)
Revised: where inequality ( i ) holds because of (1)

Original: [Equation]
Revised: [Equation]

Original: and inequality ( ii ) holds by setting ρ = 0 .
Revised: and inequality ( ii ) holds by setting ρ = 0 .

Original: Now it remains to bound
Revised: Now it remains to bound

Original: [Equation]
Revised: [Equation]

Original: using the difference ∥
Revised: using the difference ∥

Original: ω k +1 − ω k
Revised: ω k +1 − ω k

Original: ∥ 2 .
Revised: ∥ 2 .

Original: Let P
Revised: Let P

Original: kh
Revised: kh

Original: =
Revised: =

Original: P π
Revised: P π

Original: ⋆ωk h
Revised: ⋆ωk h

Original: ( s ′ | s 0 ∼ µ k ) denote the state visitation distribution resulting from π ⋆ω k probability starting at µ k
Revised: ( s ′ | s 0 ∼ µ k ) denote the state visitation distribution resulting from π ⋆ω k probability starting at µ k

Original: [Equation]
Revised: [Equation]

Original: then we have
Revised: then we have

Original: [Equation]
Revised: [Equation]

Original: Taking absolute value on both side, yields
Revised: Taking absolute value on both side, yields

Original: [Equation]
Revised: [Equation]

Original: where inequality ( i ) holds by applying Lemma B.2 with c 

Revised: where inequality ( i ) holds by applying Lemma B.2 with c

Original: = L r /α (1 − γ ) from and equality ( ii ) holds because the initial distribution of P kh is µ k .
Revised: = L r /α (1 − γ ) from and equality ( ii ) holds because the initial distribution of P kh is µ k .

Original: By the definition of d πµ , we have d π ⋆ωk
Revised: By the definition of d πµ , we have d π ⋆ωk

Original: [Equation]
Revised: [Equation]

Original: where in equality ( i )
Revised: where in equality ( i )

Original: [Equation]
Revised: [Equation]

Original: we use d k to denote d π ⋆ωk µ k .
Revised: we use d k to denote d π ⋆ωk µ k .

Original: Adding ℓ 1 norm on both sides of (14) and applying (13), yields
Revised: Adding ℓ 1 norm on both sides of (14) and applying (13), yields

Original: [Equation]
Revised: [Equation]

Original: where equality ( i ) holds because (cid:80) ∞
Revised: where equality ( i ) holds because (cid:80) ∞

Original: h =0 γ h h =
Revised: h =0 γ h h =

Original: γ/ (1 − γ ) 2 and equality ( ii ) holds because of (1).
Revised: γ/ (1 − γ ) 2 and equality ( ii ) holds because of (1).

Original: Hence
Revised: Hence

Original: [Equation]
Revised: [Equation]

Original: we know that
Revised: we know that

Original: [Equation]
Revised: [Equation]

Original: where ∆ kω = max 1 ≤ i ≤ k ∥ ω
Revised: where ∆ kω = max 1 ≤ i ≤ k ∥ ω

Original: i − ω
Revised: i − ω

Original: i − 1 ∥ 2 and the last ≈ holds when k is large.
Revised: i − 1 ∥ 2 and the last ≈ holds when k is large.

Original: Therefore, applying (16) back to (11), we know that
Revised: Therefore, applying (16) back to (11), we know that

Original: [Equation]
Revised: [Equation]

Original: where inequality ( i ) holds since Lemma B.2 implies c
Revised: where inequality ( i ) holds since Lemma B.2 implies c

Original: 1 = L r /α (1 − γ ) , and we omit the 1 / (1 − γ )and log in the (cid:101) Ω , which completes the proof.
Revised: 1 = L r /α (1 − γ ) , and we omit the 1 / (1 − γ )and log in the (cid:101) Ω , which completes the proof.

Original: Note that we can only achieve the final bound (cid:101) Ω (cid:16) L r α (1 − β ) ∆ kω S (cid:17) by setting β as a constant.
Revised: Note that we can only achieve the final bound (cid:101) Ω (cid:16) L r α (1 − β ) ∆ kω S (cid:17) by setting β as a constant.

Original: If we pick an arbitrarily small β , then the 1 /β term will dominate the complexity and we will not have the final bound of (cid:101) Ω (cid:16)
Revised: If we pick an arbitrarily small β , then the 1 /β term will dominate the complexity and we will not have the final bound of (cid:101) Ω (cid:16)

Original: L r α (1 − β ) ∆ kω S (cid:17) .
Revised: L r α (1 − β ) ∆ kω S (cid:17) .

Original: ■
Revised: ■

Original: A. C OMPLEXITY OF V ANILLA S TOCHASTIC PG
Revised: A. C OMPLEXITY OF V ANILLA S TOCHASTIC PG

Original: Theorem A.2 (Complexity of Stochastic PG (Theorem 5.1 of Ding et al. (2021))) Consider an arbitrary tolerance level δ > 0 and a small enough tolerance level ε > 0 .
Revised: Theorem A.2 (Complexity of Stochastic PG (Theorem 5.1 of Ding et al. (2021))) Consider an arbitrary tolerance level δ > 0 and a small enough tolerance level ε > 0 .

Original: For every initial point θ 1 , if θ T +1 is generated by SPG (Algorithm 4) with
Revised: For every initial point θ 1 , if θ T +1 is generated by SPG (Algorithm 4) with

Original: [Equation]
Revised: [Equation]

Original: where
Revised: where

Original: [Equation]
Revised: [Equation]

Original: then we have P ( D ( θ T +1 ≤ ε )) ≥
Revised: then we have P ( D ( θ T +1 ≤ ε )) ≥

Original: 1 − δ .
Revised: 1 − δ .

Original: Corollary A.3 (Iteration Complexity and Sample Complexity for ε -Optimal Policies) Suppose we set the tolerance level ε, δ
Revised: Corollary A.3 (Iteration Complexity and Sample Complexity for ε -Optimal Policies) Suppose we set the tolerance level ε, δ

Original: = O ( S − 1 ) , the iteration complexity and sample complexity of obtaining an ε -optimal policy using stochastic softmax policy gradient (Algorithm 4) in phase 1 and phase 2 satisfies:
Revised: = O ( S − 1 ) , the iteration complexity and sample complexity of obtaining an ε -optimal policy using stochastic softmax policy gradient (Algorithm 4) in phase 1 and phase 2 satisfies:

Original: • Phase 1: T 1 = (cid:101) Ω (cid:16) S 2 S 3 (cid:17) , B 1 = (cid:101) Ω (cid:16) S 2 S 3 (cid:17) , • Phase 2: T 2 = (cid:101)
Revised: • Phase 1: T 1 = (cid:101) Ω (cid:16) S 2 S 3 (cid:17) , B 1 = (cid:101) Ω (cid:16) S 2 S 3 (cid:17) , • Phase 2: T 2 = (cid:101)

Original: Ω (cid:0) S 3 / 2 (cid:1) , B 2 = (cid:101) Ω (cid:0) S 5 (cid:1) ,
Revised: Ω (cid:0) S 3 / 2 (cid:1) , B 2 = (cid:101) Ω (cid:0) S 5 (cid:1) ,

Original: with probability at least 1 − δ .
Revised: with probability at least 1 − δ .

Original: Proof We first check the dependency of (19) on S .
Revised: Proof We first check the dependency of (19) on S .

Original: Notice that
Revised: Notice that

Original: [Equation]
Revised: [Equation]

Original: Hence, the complexities in phase 1 scales at
Revised: Hence, the complexities in phase 1 scales at

Original: [Equation]
Revised: [Equation]

Original: To enforce a positive T 2 , the tolerance level ε, δ should satisfy ε 0 6 δε ≥ 1 , which implies 1 δε = Ω( S 2 ) .
Revised: To enforce a positive T 2 , the tolerance level ε, δ should satisfy ε 0 6 δε ≥ 1 , which implies 1 δε = Ω( S 2 ) .

Original: Hence, if assuming ε 0 δε = o ( S ) the tolerance level ε, δ
Revised: Hence, if assuming ε 0 δε = o ( S ) the tolerance level ε, δ

Original: = O ( S − 1 ) , the complexities in phase 2 scales at
Revised: = O ( S − 1 ) , the complexities in phase 2 scales at

Original: T2 ≥
Revised: T

Original: [Equation]
Revised: [Equation]

Original: A. C OMPLEXITY OF L EARNING THE N EXT C ONTEXT
Revised: A. C OMPLEXITY OF L EARNING THE N EXT C ONTEXT

Original: Theorem A.4 (Theorem 4.1: Complexity of Learning the Next Context) Consider the context based stochastic softmax policy gradient (line 6 of Algorithm 1), if ∥ ω k − ω k − 1 ∥ 2 ≤ O (cid:0) S − 2 (cid:1) , andthe tolerance level ε, δ
Revised: Theorem A.4 (Theorem 4.1: Complexity of Learning the Next Context) Consider the contextbased stochastic softmax policy gradient (line 7 of Algorithm 1), suppose Assumption 3.1 and

Original: = O ( S − 1 ) , then the iteration complexity of obtaining an ε -optimal policy for ω k froman θ ⋆ω k − 1 is  Ω (cid:0) S 3 / 2 (cid:1) and the per iteration sample complexity is (cid:101) Ω (cid:16)
Revised: Assumption 3.2 hold, then the iteration number of obtaining an ε -optimal policy for ω k from θ ⋆ω k − 1 is (cid:101) Ω (cid:0) S 3 / 2 (cid:1) and the per iteration sample complexity is (cid:101) Ω (cid:16)

Original: L r α (1 − β ) S 3 (cid:17) .
Revised: L r α (1 − β ) S 3 (cid:17) .

Original: We first introduce the following lemma to aid the proof of Theorem A.4.
Revised: We first introduce the following lemma to aid the proof of Theorem A.4.

Original: Lemma A.5 (Bounded Optimal Values Between two Adjacent Contexts)
Revised: Lemma A.5 (Bounded Optimal Values Between two Adjacent Contexts)

Original: Under the same conditions as Theorem A.4, we have
Revised: Under the same conditions as Theorem A.4, we have

Original: [Equation]

Original: [Equation]
Revised: [Equation]

Original: Proof Let V πω denote the value function of policy π with reward function r ω .
Revised: Proof Let V πω denote the value function of policy π with reward function r ω .

Original: From (66) of
Revised: From (65) of

Original: Lemma B.3, we know that for any initial distribution ρ , we have
Revised: Lemma B.3, we know that for any initial distribution ρ , we have

Original: [Equation]
Revised: [Equation]

Original: From (48) of Lemma B.1, we know that
Revised: From (47) of Lemma B.1, we know that

Original: [Equation]
Revised: [Equation]

Original: hence
Revised: hence

Original: [Equation]
Revised: [Equation]

Original: we have
Revised: we have

Original: [Equation]
Revised: [Equation]

Original: Let f ( x ) denote the log soft max function for an input vector x =
Revised: Let f ( x ) denote the log soft max function for an input vector x =

Original: [ x 1 , x 2 , . . .
Revised: [ x 1 , x 2 , . . .

Original: , x A ] ⊤ such that x i ≥ 0 , then for a small perturbation ∆ ∈ R A , the intermediate value theorem implies
Revised: , x A ] ⊤ such that x i ≥ 0 , then for a small perturbation ∆ ∈ R A , the intermediate value theorem implies

Original: [Equation]
Revised: [Equation]

Original: for some vector z on the segment [ x , x + ∆ ] .
Revised: for some vector z on the segment [ x , x + ∆ ] .

Original: Now consider the Jacobian of the log softmax function ∂ [ ∇ z f ( z )]
Revised: Now consider the Jacobian of the log softmax function ∂ [ ∇ z f ( z )]

Original: i /∂z j :
Revised: i /∂z j :

Original: [Equation] [Equation]
Revised: [Equation] [Equation]

Original: where p i ( z ) = exp( z i ) / (cid:80) Ak =1 exp( z k ) .
Revised: where p i ( z ) = exp( z i ) / (cid:80) Ak =1 exp( z k ) .

Original: hence, we know that
Revised: hence, we know that

Original: [Equation]
Revised: [Equation]

Original: Now let x =
Revised: Now let x =

Original: [Equation]
Revised: [Equation]

Original: 33) (58) from Lemma B.2 implies that
Revised: 33) (57) from Lemma B.2 implies that

Original: [Equation]
Revised: [Equation]

Original: substituting (34) and (32) into (29), yields
Revised: substituting (34) and (32) into (29), yields

Original: [Equation]
Revised: [Equation]

Original: Combine (35) with (27)
Revised: Combine (35) with (27)

Original: [Equation]
Revised: [Equation]

Original: we have
Revised: we have

Original: [Equation]
Revised: [Equation]

Original: which completes the proof
Revised: which completes the proof

Original: [Equation]
Revised: [Equation]

Original: ■
Revised: ■

Original: Now we are ready to proceed to the proof of Theorem A.4.
Revised: Now we are ready to proceed to the proof of Theorem A.4.

Original: Proof From (19) we know that
Revised: Proof From (19) we know that

Original: [Equation]
Revised: [Equation]

Original: And from Section 6.2 of Ding et al.
Revised: And from Section 6.2 of Ding et al.

Original: [Equation]
Revised: [Equation]

Original: we can directly enter phase 2 of the stochastic PG
Revised: we can directly enter phase 2 of the stochastic PG

Original: when
Revised: when

Original: [Equation] [Equation]
Revised: [Equation] [Equation]

Original: Hence, when ∆ kω = max 1 ≤ i ≤ k ∥ ω
Revised: Hence, when ∆ kω = max 1 ≤ i ≤ k ∥ ω

Original: i − ω 
Revised: i − ω i

Original: i − 1 ∥ 2 = O (1 /S 2 ) , we have
Revised: [Equation] 1 ∥ 2 = O (1 /S 2 ) , we have

Original: [Equation]
Revised: [Equation]

Original: which implies we can directly enter phase 2 and enjoys the faster iteration complexity of T 

Revised: which implies we can directly enter phase 2 and enjoys the faster iteration complexity of T

Original: = Ω (cid:0) S 3 / 2 (cid:1) (by choosing δ
Revised: = Ω (cid:0) S 3 / 2 (cid:1) (by choosing δ

Original: = O ( S − 1 ) ) and the smaller batch size of
Revised: = O ( S − 1 ) ) and the smaller batch size of

Original: [Equation]
Revised: [Equation]

Original: where equation ( i ) holds by applying Lemma A.1 to (23):
Revised: where equation ( i ) holds by applying Lemma A.1 to (23):

Original: [Equation]
Revised: [Equation]

Original: and equality ( ii ) holds by the assumption that ∆ kω =
Revised: and equality ( ii ) holds by the assumption that ∆ kω =

Original: O ( S − 2 )
Revised: O ( S − 2 )

Original: and we omit the log term and components not related to S in (cid:101) Ω .
Revised: and we omit the log term and components not related to S in (cid:101) Ω .

Original: ■
Revised: ■

Original: A. T OTAL C OMPLEXITY OF R OLLIN
Revised: A. T OTAL C OMPLEXITY OF R OLLIN

Original: Theorem A.6 (Theorem 4.2: Total Complexity of Learning the Target Context) Supposethe assumptions in Theorem 4.1 holds.
Revised: Theorem A.6 (Theorem 4.3: Total Complexity of Learning the Target Context) Suppose

Original: Given a curriculum { ω k } Kk =0 such that max [Equation] ) , the total number of iteration of learning π ⋆ω K using Algorithm 1 is Ω( KS 3 / 2 ) and the per iteration is (cid:101) Ω (cid:0) S 3 (cid:1) , with high probability.
Revised: Assumption 3.1 and 3.2 hold, and θ (0) is an near-optimal initialization, then the total number of iteration of learning π ⋆ω K using Algorithm 1 is Ω( KS 3 / 2 ) and the per iteration is (cid:101) Ω (cid:0) S 3 (cid:1) , with high probability.

Original: Proof From lemma A.5, we know that
Revised: Proof From lemma A.5, we know that

Original: [Equation]
Revised: [Equation]

Original: Suppose for each context ω k , we initialize the parameters of the policy as θ ( k )
Revised: Suppose for each context ω k , we initialize the parameters of the policy as θ ( k )

Original: =
Revised: =

Original: θ ⋆ω k − 1 , and let θ ( k ) t denote the parameters at the t th iteration of SPG.
Revised: θ ⋆ω k − 1 , and let θ ( k ) t denote the parameters at the t th iteration of SPG.

Original: We will use induction to show that when t = (cid:101) Ω( S 3 / 2 ) , ∀ k ∈
Revised: We will use induction to show that when t = (cid:101) Ω( S 3 / 2 ) , ∀ k ∈

Original: [ K ] , we have
Revised: [ K ] , we have

Original: [Equation]
Revised: [Equation]

Original: (42) this implies that for any context ω k , k ∈ [ K ] , we can always find a good initialization by setting θ ( k )
Revised: (42) this implies that for any context ω k , k ∈ [ K ] , we can always find a good initialization by setting θ ( k )

Original: = θ ( k − 1) t from learning π ⋆ω k − 1 using SPG after t =
Revised: = θ ( k − 1) t from learning π ⋆ω k − 1 using SPG after t =

Original: Ω( S 3 / 2 ) iteration.
Revised: Ω( S 3 / 2 ) iteration.

Original: This result guarantees that every initialization θ ( k ) 1 for learning the optimal contextual policy π ⋆ω k will directly start from the efficient phase 2.
Revised: This result guarantees that every initialization θ ( k ) 1 for learning the optimal contextual policy π ⋆ω k will directly start from the efficient phase 2.

Original: Induction: k = 1 .
Revised: Induction: k = 0 .

Original: When k = 1 , from (39) of Theorem A.4, we know thatat k = 1 , we have
Revised: When k = 0 , Assumption 3.2 and the near-optimal initialization (Definition 4.2) of θ (0)0 implies that

Original: [Equation]
Revised: [Equation]

Original: Hence, if we initialize θ (1)1 such that θ (1)1 = θ ⋆ω 0 , then we have:

Original: [Equation]
Revised: [Equation]

Revised: This result implies that a near-optimal initialization allows the initialization to directly start from phase 2 of SPG.

Original: Induction: from k − 1 to k .
Revised: Induction: from k − 1 to k .

Original: Suppose the result in (42) holds for k − 1 , then we know that
Revised: Suppose the result in (42) holds for k − 1 , then we know that

Original: [Equation]
Revised: [Equation]

Original: Select ε such that ε ≤ ε 0 / 2 .
Revised: Select ε such that ε ≤ ε 0 / 2 .

Original: Theorem A.4 suggests that when t ′ = ˜Ω( S 3 / 2 ) , with high probability, we have
Revised: Theorem A.4 suggests that when t ′ = ˜Ω( S 3 / 2 ) , with high probability, we have

Original: [Equation] [Equation]
Revised: [Equation] [Equation]

Original: Hence, if we initialize θ ( k )=θ
Revised: Hence, if we initialize θ ( k )=

Original: [Equation]
Revised: [Equation]

Original: ) t , with high probability when t ′ = (cid:101) Ω( S 3 / 2 ) , we have
Revised: ) t , with high probability when t ′ = (cid:101) Ω( S 3 / 2 ) , we have

Original: [Equation]
Revised: [Equation]

Original: where inequality ( i ) holds by equation (39) in Theorem A.4, inequality ( ii ) holds because of the induction assumption in (46).
Revised: where inequality ( i ) holds by equation (39) in Theorem A.4, inequality ( ii ) holds because of the induction assumption in (45).

Original: Therefore, we have shown (42) holds for t = (cid:101) Ω( S 3 / 2 ) , ∀ k ∈
Revised: Therefore, we have shown (42) holds for t = (cid:101) Ω( S 3 / 2 ) , ∀ k ∈

Original: [ K ] .
Revised: [ K ] .

Original: Since we have K contexts in total, we know that Algorithm 1 can enforce a good initialization θ ( k )that directly starts from phase 2 for learning all π ⋆ω k , and for each k ∈
Revised: Since we have K contexts in total, we know that Algorithm 1 can enforce a good initialization θ ( k ) 1 that directly starts from phase 2 for learning all π ⋆ω k , and for each k ∈

Original: [ K ] , the iteration complexity is (cid:101) Ω( S 3 / 2 ) .
Revised: [ K ] , the iteration complexity is (cid:101) Ω( S 3 / 2 ) .

Original: Hence the total iteration complexity of obtaining an ε -optimal policy for the final context ω K is (cid:101) Ω (cid:0) KS 3 / 2 (cid:1) , with per iteration sample complexity of (cid:101) Ω (cid:0) S 3 (cid:1) .
Revised: Hence the total iteration complexity of obtaining an ε -optimal policy for the final context ω K is (cid:101) Ω (cid:0) KS 3 / 2 (cid:1) , with per iteration sample complexity of (cid:101) Ω (cid:0) S 3 (cid:1) .

Original: ■
Revised: ■

Original: A PPENDIX B K EY L EMMAS B. O PTIMAL P OLICY OF M AXIMUM E NTROPY RL N ACHUM ET AL . (2017)
Revised: A PPENDIX B K EY L EMMAS B. O PTIMAL P OLICY OF M AXIMUM E NTROPY RL N ACHUM ET AL . (2017)

Original: Lemma B.1 The optimal policy π ⋆ that maximizes the α -MaxEnt RL objective (4) with penalty term α satisfies:
Revised: Lemma B.1 The optimal policy π ⋆ that maximizes the α -MaxEnt RL objective (4) with penalty term α satisfies:

Original: [Equation]
Revised: [Equation]

Original: for all h ∈ N , where
Revised: for all h ∈ N , where

Original: [Equation] [Equation]
Revised: [Equation] [Equation]

Original: Proof Similar proof appears in (Nachum et al., 2017), we provide the proof for completeness.
Revised: Proof Similar proof appears in (Nachum et al., 2017), we provide the proof for completeness.

Original: At the optimal policy π θ = π ⋆ , take the gradient of (4) w.r.t.
Revised: At the optimal policy π θ = π ⋆ , take the gradient of (4) w.r.t.

Original: p ∈ ∆( A ) and set it to 0, we have ∂
Revised: p ∈ ∆( A ) and set it to 0, we have [Equation]

Original: [Equation]
Revised: [Equation]

Original: which implies
Revised: which implies

Original: [Equation]
Revised: [Equation]

Original:  exp (cid:18) Q π ⋆ ( s, a ) α (cid:19) .
Revised: ∝ exp (cid:18) [Equation] π ⋆ ( [Equation]

Original:  Hence, we conclude that π ⋆ ( a | s ) ∝ exp( Q ⋆ ( s, a ) /α ) .
Revised: a ) α [Equation] Hence, we conclude that π ⋆ ( a | s ) ∝ exp( Q ⋆ ( s, a ) /α ) .

Original: ■
Revised: ■

Original: B. B OUNDING THE D IFFERENCE BETWEEN O PTIMAL P OLICIES
Revised: B. B OUNDING THE D IFFERENCE BETWEEN O PTIMAL P OLICIES

Original: Lemma B.2 Suppose Assumption 3.1 holds, let π ⋆ω ( a | s ) , π ⋆ω ′ ( a | s ) denote the optimal policy for α -MaxEnt RL (48) , then ∀ ( s, a ) ∈ S × A , the optimal policies of α -MaxEnt RL under context ω, ω ′ satisfy:[Equation]
Revised: Lemma B.2 Suppose Assumption 3.1 holds, let π ⋆ω ( a | s ) , π ⋆ω ′ ( a | s ) denote the optimal policy for α -MaxEnt RL (47) , then ∀ ( s, a ) ∈ S × A , the optimal policies of α -MaxEnt RL under context ω, ω ′ satisfy:

Original: [Equation]
Revised: [Equation]

Original: Proof From Lemma C.1, we know that the soft value iteration
Revised: Proof From Lemma C.1, we know that the soft value iteration

Original: [Equation]
Revised: [Equation]

Original: is a contraction.
Revised: is a contraction.

Original: Let Q tω , Q tω ′
Revised: Let Q tω , Q tω ′

Original: denote the Q functions at the t th value iteration under context ω, ω ′ respectively
Revised: denote the Q functions at the t th value iteration under context ω, ω ′ respectively

Original: , we know Q ∞ ω =
Revised: , we know Q ∞ ω =

Original: Q π ⋆ ω and Q ∞ ω ′ =
Revised: Q π ⋆ ω and Q ∞ ω ′ =

Original: Q π ⋆ ω ′ .
Revised: Q π ⋆ ω ′ .

Original: Let ε t =
Revised: Let ε t =

Original: ∥
Revised: ∥

Original: Q tω − Q tω ′ ∥ ∞ , then we have
Revised: Q tω − Q tω ′ ∥ ∞ , then we have

Original: [Equation]

Original: [Equation]
Revised: [Equation]

Original:  where the last inequality holds because f ( x ) = log (cid:80) ni =1 exp( x i ) is a contraction.
Revised: 53) where the last inequality holds because f ( x ) = log (cid:80) ni =1 exp( x i ) is a contraction.

Original: From (54), we have
Revised: From (53), we have

Original: [Equation]
Revised: [Equation]

Original: which implies
Revised: which implies

Original: [Equation]
Revised: [Equation]

Original: where the last inequality holds by Assumption 3.1.
Revised: where the last inequality holds by Assumption 3.1.

Original: Hence, we have
Revised: Hence, we have

Original: [Equation] [Equation]
Revised: [Equation] [Equation]

Original: which implies
Revised: which implies

Original: [Equation]
Revised: [Equation]

Original: Next
Revised: Next

Original: [Equation]
Revised: [Equation]

Original: let π ⋆ω
Revised: let π ⋆ω

Original: [Equation]
Revised: [Equation]

Original: π ⋆ω ′ denote the maximum entropy policy RL under context ω
Revised: π ⋆ω ′ denote the maximum entropy policy RL under context ω

Original: [Equation]
Revised: [Equation]

Original: ω ′ respectively.
Revised: ω ′ respectively.

Original: Then for a fixed state action pair ( s, a ) ∈ S × A , we have
Revised: Then for a fixed state action pair ( s, a ) ∈ S × A , we have

Original: [Equation]
Revised: [Equation]

Original: where Q π ⋆ω ( · , s ) , Q π ⋆ω ′ ( · , s ) ∈ R
Revised: where Q π ⋆ω ( · , s ) , Q π ⋆ω ′ ( · , s ) ∈ R

Original: A
Revised: A

Original: [Equation]
Revised: [Equation]

Original: and we want to bound | π ⋆ω ( a | s )
Revised: and we want to bound | π ⋆ω ( a | s )

Original: − π ⋆ω ′ ( a | s ) | .
Revised: − π ⋆ω ′ ( a | s ) | .

Original: Next we will use (58) to bound | π ⋆ω ( a | s ) − π ⋆ω ′ ( a | s ) | , where the last inequality holds by (57).
Revised: Next we will use (57) to bound | π ⋆ω ( a | s ) − π ⋆ω ′ ( a | s ) | , where the last inequality holds by (56).

Original: Let f ( x ) denote the softmax function for an input vector x
Revised: Let f ( x ) denote the softmax function for an input vector x

Original: = [ x 1 , x 2 , . . .
Revised: = [ x 1 , x 2 , . . .

Original: , x A ] ⊤ such that x i ≥ 0 , then for a small perturbation ∆ ∈ R A , the intermediate value theorem implies[Equation]
Revised: , x A ] ⊤ such that x i ≥ 0 , then for a small perturbation ∆ ∈ R A , the intermediate value theorem implies

Original: [Equation]
Revised: [Equation]

Original:  for some vector z on the segment [ x
Revised: ) for some vector z on the segment [ x

Original: [Equation]
Revised: [Equation]

Original: x + ∆ ] .
Revised: x + ∆ ] .

Original: Hence
Revised: Hence

Original: [Equation]
Revised: [Equation]

Original: ≤ ∥ ∆ ∥
Revised: ≤ ∥ ∆ ∥

Original: [Equation]
Revised: [Equation]

Original: where the Jacobian of the softmax function ∂ [ ∇ x f ( z )] i /∂z j satisfies:
Revised: where the Jacobian of the softmax function ∂ [ ∇ x f ( z )] i /∂z j satisfies:

Original: [Equation]
Revised: [Equation]

Original: otherwise ,
Revised: otherwise ,

Original: [Equation]
Revised: [Equation]

Original: 62) and p i ( z ) =
Revised: 61) and p i ( z ) =

Original: exp( z i ) / (cid:80) Ak =1 exp( z k ) .
Revised: exp( z i ) / (cid:80) Ak =1 exp( z k ) .

Original: Now let
Revised: Now let

Original: [Equation]
Revised: [Equation]

Original: We know that f ( x ) =
Revised: We know that f ( x ) =

Original: π ⋆ω ( a | s ) and f ( x + ∆ ) =
Revised: π ⋆ω ( a | s ) and f ( x + ∆ ) =

Original: π ⋆ω ′ ( a | s ) .
Revised: π ⋆ω ′ ( a | s ) .

Original: Then (58) implies that
Revised: Then (57) implies that

Original: [Equation]
Revised: [Equation]

Original: substituting this bound on ∥ ∆ ∥ ∞ into (61), we have
Revised: substituting this bound on ∥ ∆ ∥ ∞ into (60), we have

Original: [Equation]
Revised: [Equation]

Original: which completes the proof.
Revised: which completes the proof.

Original: ■
Revised: ■

Original: B. S
Revised: B. S

Original: [Equation]
Revised: [Equation]

Original: Lemma B.

Revised: Lemma B.

Original: For any policy π and any initial distribution ρ , the value function V π ( ρ ) of the αMaxEnt RL (49) satisfies:
Revised: For any policy π and any initial distribution ρ , the value function V π ( ρ ) of the αMaxEnt RL (48) satisfies:

Original: [Equation]
Revised: [Equation]

Original: where π ⋆ is the optimal policy of the α -MaxEnt RL (4) .
Revised: where π ⋆ is the optimal policy of the α -MaxEnt RL (4) .

Original: Proof Similar proof appears in Lemma 25 & 26 of Mei et al.
Revised: Proof Similar proof appears in Lemma 25 & 26 of Mei et al.

Original: (2020), we provide the proof here for completeness.
Revised: (2020), we provide the proof here for completeness.

Original: Soft performance difference.
Revised: Soft performance difference.

Original: We first show a soft performance difference result for the MaxEnt value function (Lemma 25 of Mei et al.
Revised: We first show a soft performance difference result for the MaxEnt value function (Lemma 25 of Mei et al.

Original: By the definition of MaxEnt value function and
Revised: By the definition of MaxEnt value function and

Original: Q -function (4), (6), ∀ π, π ′ , we have
Revised: Q -function (4), (6), ∀ π, π ′ , we have

Original: [Equation]
Revised: [Equation]

Original: =
Revised: =

Original: [Equation]
Revised: [Equation]

Original: where the last equality holds because by the definition of state visitation distribution
Revised: where the last equality holds because by the definition of state visitation distribution

Original: [Equation]
Revised: [Equation]

Original: taking expectation of s w.r.t.
Revised: taking expectation of s w.r.t.

Original: s ∼ ρ , yields
Revised: s ∼ ρ , yields

Original: [Equation]
Revised: [Equation]

Original: and (69) is known as the soft performance difference lemma (Lemma 25 in Mei et al.
Revised: and (68) is known as the soft performance difference lemma (Lemma 25 in Mei et al.

Original: [Equation]
Revised: [Equation]

Original: Soft sub-
Revised: Soft sub-

Original: optimality.
Revised: optimality.

Original: Next we will show the soft sub-optimality result.
Revised: Next we will show the soft sub-optimality result.

Original: By the definition of the optimal policy of α -MaxEnt RL (48), we have
Revised: By the definition of the optimal policy of α -MaxEnt RL (47), we have

Original: [Equation]
Revised: [Equation]

Original: Substituting π ⋆ into the performance difference lemma (69), we have
Revised: Substituting π ⋆ into the performance difference lemma (68), we have

Original: [Equation]
Revised: [Equation]

Original: taking expectation s ∼ ρ yields
Revised: taking expectation s ∼ ρ yields

Original: [Equation]
Revised: [Equation]

Original: which completes the proof
Revised: which completes the proof

Original: [Equation]
Revised: [Equation]

Original: ■
Revised: ■

Original: A PPENDIX C S UPPORTING L EMMAS C. B ELLMAN C ONSISTENCY E QUATION OF M AX E NT RL
Revised: A PPENDIX C S UPPORTING L EMMAS C. B ELLMAN C ONSISTENCY E QUATION OF M AX E NT RL

Original: Lemma C.1 (Contraction of Soft Value Iteration) From (49) and (6) , the soft value iteration operator T defined as
Revised: Lemma C.1 (Contraction of Soft Value Iteration) From (48) and (6) , the soft value iteration operator T defined as

Original: [Equation]
Revised: [Equation]

Original: is a contraction.
Revised: is a contraction.

Original: Proof A similar proof appears in Haarnoja (2018), we provide the proof for completeness.
Revised: Proof A similar proof appears in Haarnoja (2018), we provide the proof for completeness.

Original: To see (73) is a contraction, for each ( s, a
Revised: To see (72) is a contraction, for each ( s, a

Original: ) ∈ S × A , we have
Revised: ) ∈ S × A , we have

Original: [Equation]
Revised: [Equation]

Original: which implies T Q 1 ( s, a ) − T Q 2 ( s, a ) ≤ γ ∥ Q 1 − Q 2 ∥ ∞ .
Revised: which implies T Q 1 ( s, a ) − T Q 2 ( s, a ) ≤ γ ∥ Q 1 − Q 2 ∥ ∞ .

Original: Similarly, we also have T Q 2 ( s, a ) −
Revised: Similarly, we also have T Q 2 ( s, a ) −

Original: T Q 1 ( s, a ) ≤ γ ∥ Q 1 − Q 2 ∥ ∞ , hence we conclude that
Revised: T Q 1 ( s, a ) ≤ γ ∥ Q 1 − Q 2 ∥ ∞ , hence we conclude that

Original: [Equation]
Revised: [Equation]

Original: which implies ∥ Q 1 − Q 2 ∥ ∞ ≤ γ ∥T Q 1 − T Q 2 ∥ ∞ .
Revised: which implies ∥ Q 1 − Q 2 ∥ ∞ ≤ γ ∥T Q 1 − T Q 2 ∥ ∞ .

Original: Hence T is a γ -contraction and the optimal policy π ⋆ of it is unique.
Revised: Hence T is a γ -contraction and the optimal policy π ⋆ of it is unique.

Original: ■
Revised: ■

Original: C. C ONSTANT M INIMUM P OLICY P ROBABILITY
Revised: C. C ONSTANT M INIMUM P OLICY P ROBABILITY

Original: Lemma C.2 (Lemma 16 of Mei et al.
Revised: Lemma C.2 (Lemma 16 of Mei et al.

Original: Using the policy gradient method (Algorithm 3) with an initial distribution ρ such that ρ ( s ) > 0 , ∀ S , we have
Revised: Using the policy gradient method (Algorithm 3) with an initial distribution ρ such that ρ ( s ) > 0 , ∀ S , we have

Original: [Equation]
Revised: [Equation]

Original: is a constant that does not dependent on t .
Revised: is a constant that does not dependent on t .

Original: Remark C.3 (State Space Dependency of constant c ) Note that assuming c is independent of S is quite strong and does not generally hold as suggested by Li et al.
Revised: Remark C.3 (State Space Dependency of constant c ) Note that assuming c is independent of S is quite strong and does not generally hold as suggested by Li et al.

Original: Still, if one replace the constant c with other S dependent function f ( S ) , one still can apply a similar proof technique for
Revised: Still, if one replace the constant c with other S dependent function f ( S ) , one still can apply a similar proof technique for

Original: Theorem 4.1 to show that R OLLIN reduces the iteration complexity, and the final iteration complexity bound in Theorem 4.1 will include an additional f ( S ) .
Revised: Theorem 4.1 to show that R OLLIN reduces the iteration complexity, and the final iteration complexity bound in Theorem 4.1 will include an additional f ( S ) .

Original: A PPENDIX D S UPPORTING A LGORITHMS
Revised: A PPENDIX D S UPPORTING A LGORITHMS

Original: Algorithm 3 PG for α -MaxEnt RL (Algorithm 1 in Mei et al.
Revised: Algorithm 3 PG for α -MaxEnt RL (Algorithm 1 in Mei et al.

Original: Algorithm
Revised: Algorithm

Original: ▷ Phase
Revised: ▷ Phase

Original: ▷ Phase7:
Revised: ▷ Phase7:

Original: ▷ Algorithm9: 
Revised: ▷ Algorithm9: end

Original: 1: Input: ρ, α, θ 0 , B, t, η t 2: for i = 1 , 2 , ..., B do 3: s iH
Revised: 1: Input: ρ, α, θ 0 , B, t, η t 2: for i = 1 , 2 , ..., B do 3: s iH

Original: [Equation]
Revised: [Equation]

Original: a
Revised: a

Original: [Equation]
Revised: [Equation]

Original: ← SamSA
Revised: ← SamSA

Original: [Equation]
Revised: [Equation]

Original: ρ,
Revised: ρ,

Original: [Equation] [Equation]
Revised: [Equation] [Equation]

Original: ▷ Algorithm4:
Revised: ▷ Algorithm4:

Original: ˆ Q
Revised: ˆ Q

Original: [Equation]
Revised: [Equation]

Original: ← EstEntQ
Revised: ← EstEntQ

Original: [Equation]
Revised: [Equation]

Original: ▷ Algorithm5: end for 6: θ
Revised: ▷ Algorithm5: end for 6: θ

Original: [Equation]
Revised: [Equation]

Original: ←
Revised: ←

Original: [Equation]
Revised: [Equation]

Original: Remark D.1 Lemma 3.4 in Ding et al.
Revised: Remark D.1 Lemma 3.4 in Ding et al.

Original: (2021) implies that the estimator
Revised: (2021) implies that the estimator

Original: [Equation]
Revised: [Equation]

Original: (77) in line 6 of Algorithm 6 is an unbiased estimator of the gradient ∇ θ
Revised: in line 6 of Algorithm 6 is an unbiased estimator of the gradient ∇ θ

Original: V π θ ( ρ ) .
Revised: V π θ ( ρ ) .

Original: γ s do 5:

Original: : end
Revised: : end

Revised: H

Revised: H γ h r s a h

Revised: 1 ) − α log π θ

Original: A PPENDIX E E XPERIMENTAL D ETAILS
Revised: A PPENDIX E E XPERIMENTAL D ETAILS

Revised: We use the

Revised: SAC implementation from https://github.com/ikostrikov/ jaxrl (Kostrikov, 2021) for all our experiments in the paper.

Revised: The exception is the goal reaching task with automated curriculum generation, where we use our customized DDPG implementation that is highly inspired from the same codebase.

Revised: E. G OAL R EACHING WITH AN O RACLE C URRICULUM

Original: For our antmaze-umaze experiments, We use a sparse reward function where the reward iswhen the distance D between the ant and the goal is greater than 0 .
Revised: For our antmaze-umaze experiments with oracle curriculum, we use a sparse reward function where the reward is 0 when the distance D between the ant and the goal is greater than 0 .

Original: 5 and r = exp( − 5 D ) when the distance is smaller than or equal to 0 .
Revised: 5 and r = exp( − 5 D ) when the distance is smaller than or equal to 0 .

Original: The performance threshold is set to be R = 200 .
Revised: The performance threshold is set to be R = 200 .

Original: Exceeding such threshold means that the ant stays on top of the desired location for at least 200 out ofsteps, where 500 is the maximum episode length of the antmaze-umaze environment.
Revised: Exceeding such threshold means that the ant stays on top of the desired location for at least 200 out of 500 steps, where 500 is the maximum episode length of the antmaze-umaze environment.

Original: We use the average return of the last 10 episodes and compare it to the performance threshold R .
Revised: We use the average return of the last 10 episodes and compare it to the performance threshold R .

Original: For both of the SAC agents, we use the same set of hyperparameters shown in Table 5.
Revised: For both of the SAC agents, we use the same set of hyperparameters shown in Table 3.

Original: We use the SAC implementation from https://github.com/ikostrikov/jaxrl .

Original: See Algorithm 8, for a more detailed pseudocode.
Revised: See Algorithm 8, for a more detailed pseudocode.

Original: [Table]
Revised: [Table]

Original: [Table]

Original: then
Revised: 18: a

Original: [Table] [Table] else [Table]
Revised: 19: else 20: a main s ω

Original: [Table] [Table] if Rollin is enabled for the current episode then [Table]
Revised: − 1 ) 21: with probability

Original: if [Table] is stopped forthe [Table]
Revised: 1 − γ for

Original: episode then
Revised: 22: end

Original: [Table] [Table] else [Table]
Revised: 23: else 24: a t ∼ π main a t s t ω k ) 25: end if 26: take action a t in the environment and receives s t +1 and r t = r ω k

Original: : with probability of [Table] Rollin  for the current episode [Table]
Revised: ( s t , a t ) 27: add ( s t , a t , s t +1 , r t ) in replay buffer D 28: if Rollin is disabled for the current episode then 29: update π main using D . 30: end if 31: if π exp was used to produce a t then 32: add ( s t , a t , s t + r t ) in replay buffer D exp 33: update

Original: end if [Table] else [Table]
Revised: D

Original: 25: end if [Table] take action a t in the environment and receives s t +1 and r t = r
Revised: 34: end if 35: t ← t + 1 36: end for 37: Output: π main

Original: [Table] k ( s t [Table] a
Revised: E. G OAL R EACHING WITH AN A UTOMATICALLY G ENERATED C URRICULUM

Original: [Table] [Table] add{+For our antmaze-umaze experiments with automatically generated curriculum, we use a sparse reward function where the reward is − 1 when the distance between the ant and the goal is greater than 1 . 
+}
Revised: {+For our antmaze-umaze experiments with automatically generated curriculum, we use a sparse reward function where the reward is − 1 when the distance between the ant and the goal is greater than 1 . 
+}

Original: [Table] s t , a t , s t +1 , r t ) in replay buffer D 28: if Rollin is disabled for the current episode then 29: update π main using D main . 
Revised: (this is different from the oracle curriculum setting where the threshold is 0 . 5 ), and 0 otherwise.

Original: 30: end if 31: if π exp was used to produce a t then 32: add ( s t , a t , s t +1 , r t ) in replay buffer D exp 33: update π exp using D exp . 
Revised: We closely follow the DDPG setup in the MEGA paper (Pitis et al., 2020) and use the rfaab relabeling with a ratio of (0 . 1 , 0 . 4 , 0 . 1 , 0 . 3 , 0 . 1) . We also use the go-explore style exploration where the epsilon-greedy ratio is increased from 0 . 1 to 0 . 11 when the goal is reached (whenever a reward of zero is received).

Original: {+The maximum episode length is 
+}
Revised: {+The maximum episode length is 
+}

Revised: and we estimate the success rate of the agent by rolling out the policy 10 times with different pairs of random initial starting locations (a square region centered around (0 . 0 , 0 . 0) with length 2 . 0 ) and goals (a square region centered around (0 . 0 , − 8 . 0) with length 2 . 0 ).

Revised: See Figure 5 for a visualization of 8 successful trajectories of this task.

Revised: We use three-layer network (512 , 512 , 512) with layer normalization for both the actor and the critic.

Revised: We also use target Q value clipping such that the target Q value that gets backed up is always within [ − 100 , 0] (the range of the achievable Q value with a discount factor of 0 . 99 and binary reward {− 1 ., 0 . } ).

Revised: We did not use a dynamic normalization of the observation (which was used in the original paper).

Revised: Also, we used tanh to squash the output of the actor to be between − 1 and 1 rather than using L 2 action penalty as done in the original paper.

Revised: We include a detailed description of how R OLLIN is applied on top of MEGA in Algorithm 9 with N = 10 5 and the goal proposer from the MEGA paper with the default parameters.

Revised: See Figure 4 for the results of R OLLIN on top of MEGA over a range of β values.

Revised: [Figure]

Revised: Baseline β = 0 .β

Revised: = 0 .β

Revised: = 0 .β

Revised: = 0 .β

Revised: = 0 .

Revised: 18: a

Revised: 19: else 20: a

Revised: 22: end

Revised: 24: a

Revised: 25: end

Revised: 26: take action

Revised: 27: add

Revised: 28: if Rollin is disabled for current episode then 29:

Revised: 30: end was used to produce a t then 32:

Revised: 1 , g ) in replay buffer D exp 33: update π exp using D exp with goal relabeling.

Original: 34: end if 35: t ← t + 1 36: end for 37: Output: π main 
Revised: 34: end if 35: t ← t + 1 36: end for 37: Output: π main ,k

Revised: E. N ON G OAL R EACHING

Revised: For the non goal reaching tasks in walker2d , hopper , humanoid , and ant experiments, the desired x -velocity range [ λκ, λ ( κ + 0 . 1)) , the near-optimal threshold R ( κ ) , and the healthy_reward all depend on the environments.

Revised: The maximum episode length 1000.

Revised: Details are provided in Table 4.

Revised: [Table]

Original: A PPENDIX F A DDITIONAL L EARNING C URVES 
Revised: A PPENDIX F A DDITIONAL L EARNING C URVES AND T ABLES F. G OAL R EACHING

Original: [Figure]
Revised: [Figure]

Revised: Vanilla R OLLIN without geometric sampling

Original: [Figure] 
Revised: [Figure] [Figure]

Original:  OLLIN on an oracle curriculum in Figure 2.
Revised: on antmaze-umaze with R OLLIN on an oracle curriculum in Figure 2.

Original: The confidence interval represents the standard error computed overrandom seeds.
Revised: The confidence interval represents the standard error computed over 8 random seeds.

Original: [Table]
Revised: [Table]

Revised: L e a r n i n g

Revised: P r og r e ss ( κ ) ∆

Revised: [Table]

Revised: [Table]

Revised: L e a

Revised: [Table]

Revised: n i n g

Revised: P r og

Revised: [Table]

Revised: e ss ( κ ) ∆

Revised: [Table]

Revised: [Table]

Revised: n i n

Revised: [Table]

Revised: og

Revised: [Table]

Revised: R OLLIN + Go-Explore without geometric sampling

Original: [Table]
Revised: [Table]

Revised: L e a r n i n g

Revised: P r og r e ss ( κ ) ∆=1/

Revised: [Table]

Original: Env.
Revised: Env.

Revised: Steps (

Revised: [Table]

Revised: [Table]

Revised: L e a r n i n g

Original: Steps ( × 10 6 ) 0 .0 .1 .
Revised: P r og r e ss ( κ ) ∆=1/ [Table] [Table] [Table]

Revised: Env.

Revised: Steps ( × 10 6 )

Revised: [Table]

Revised: [Table]

Original: L e a r n i n g
Revised: L e a r n i n g

Revised: P r og r e ss ( κ ) ∆=1/

Revised: R OLLIN

Revised: [Figure]

Revised: -umaze with

Revised: R OLLIN on an oracle curriculum in Figure 2.

Revised: The confidence interval represents the standard error computed over 8 random seeds.

Revised: [Table]

Revised: Env

Revised: [Table]

Revised: [Table]

Revised: [Table]

Revised: L e a r n i n g

Original: P r og r e ss ( κ ) ∆=1/ 0 . 0 0 .
Revised: P r og r e ss ( κ ) ∆=1/ 0 . 0 0 .

Original: Env.
Revised: Env.

Original: Steps ( × 10 6 ) 0 .0 .1 .
Revised: Steps ( × 10 6 ) 0 .0 .0 .

Original: L e a r n i n g
Revised: L e a r n i n g

Original: P r og r e ss ( κ ) ∆=1/
Revised: P r og r e ss ( κ ) ∆=1/

Original: [Table]
Revised: [Table]

Original: L e a r n i n g
Revised: [Table]

Original: [Table]
Revised: [Table]

Original: r og r e 
Revised: L e a r n i n g

Original: [Table] [Table] [Table] R
Revised: P r og r e ss ( κ )

Original: [Table]
Revised: [Table]

Original:  + [Table] sampling
Revised: R OLLIN + Go-Explore without geometric sampling

Original: [Table]
Revised: [Table]

Original: Env.

Original: Steps ( [Table] [Table] [Table]
Revised: [Table]

Original: L e a r n i n g
Revised: L e a r n i n g

Original: P r og r e ss ( κ ) ∆=1/ 0 . 0 0 .
Revised: P r og r e ss ( κ ) ∆=1/ 0 . 0 0 .

Original: Env.
Revised: Env.

Original: Steps ( × 10 6 ) 0 .0.1 .
Revised: Steps ( × 10 6 ) 0 .0 .

Original: L e a r n i n g
Revised: L e a r n i n g

Original: P r og r e ss ( κ ) ∆=1/ 0 . [Table] [Table] [Table]
Revised: P r og r e ss ( κ ) ∆=1/ 0 . 0 0 .

Original: [Table]

Original: Env.
Revised: Env.

Original: Steps ( × 10 6 )0 .0 .
Revised: Steps ( × 10 6 )

Original: [Table]
Revised: [Table]

Original: L e a r n i n g
Revised: L e a r n i n g

Original: P r og r e ss ( κ ) ∆=1/
Revised: P r og r e ss ( κ ) ∆=1/

Original: R OLLIN + relabeling
Revised: R OLLIN + Go-Explore with geometric sampling

Original: [Figure]
Revised: [Table]

Original: on antmaze-umaze with R OLLIN on an oracle curriculum in Figure 2. The confidence interval represents the standard error computed over 8 random seeds.
Revised: [Table]

Original: [Table]
Revised: [Table]

Original: L e a [Table] n i n g
Revised: Table 7 Go-Explore with different exploration noise. EN represents the multiplier for the Gaussian exploration noise.

Original: P r og r e ss ( [Table] [Table] [Table]
Revised: All other settings are the same as Table 5.

Original:  R [Table] geometric sampling
Revised: F. G OAL R EACHING WITH A L EARNED C URRICULUM

Original: [Table]
Revised: [Figure]

Original: L e [Table] n i n g P r og
Revised: 2020) on antmaze-umaze with different choices of β .

Original: [Table]
Revised: [Table]

Original: e ss ( κ ) ∆=1/ [Table] L e a r n i n g P r ogr e ss ( κ ) ∆=1/
Revised: og

Original: [Table]
Revised: [Table]

Original: L e a r n i n g P r ogr e ss ( κ ) ∆=1/ R OLLIN + Go-Explore
Revised: og

Original: [Figure] [Figure] [Figure]
Revised: [Table] hopper

Original: [Figure]
Revised: [Table]

Original: Env
Revised: og

Original: [Figure]
Revised: [Table]

Revised: humanoid 0 .

Original: [Figure] [Figure]
Revised: [Table] [Table] [Table]

Original: L e a r n i n g P r og r e ss ( κ ) ∆=1/ 0 . 0 0 .
Revised: [Table] [Table] [Table]

Original: Env.
Revised: humanoid

Original: Steps ( × 10 6 ) 0 .0 .0 .
Revised: [Table] [Table] [Table] [Table]

Original: L e a r n i n g
Revised: r [Table] r n humanoid [Table]

Original: P r og r e ss ( κ ) ∆=1/ 0 . 0 0 .
Revised: [Table] [Table] [Table] [Table]

Original: Env.
Revised: og

Original: Steps ( × 10 6 ) 0 .0.0 . 
Revised: [Table] ant 0 .0 . [Table]

Original: L e a r n i n g
Revised: [Table]

Original: P r og r e ss ( κ ) ∆=1/
Revised: [Table] [Figure]

Original: Baseline β = 0 .β
Revised: [Table] [Table]

Original: = 0 .β = 0 .β = 0 .β = 0 .
Revised: [Table]

Original: R OLLIN + Go-Explore without geometric sampling [Figure] [Figure] [Figure]
Revised: Step β =

Original: [Figure]

Original: [Figure]
Revised: [Table]

Original: L e a r n i n g

Original: P r og r e ss ( κ ) ∆=1/ [Figure] [Figure] [Figure]
Revised: β =

Original: [Figure]

Original: Env

Original: [Figure] [Figure] [Figure] L e a r n i n g
Revised: [Table]

Original: P r og r e ss ( κ ) ∆=1/ [Figure] [Figure] [Figure]
Revised: β = 0 .β

Original: [Figure] [Figure] Env. Steps ( × 10 6 )
Revised: = 0 .

Original: [Figure]
Revised: [Table]

Original: a r n i n g P r og r e ss ( [Figure] ) ∆=1/
Revised: β =

Original: [Figure]
Revised: [Table]

Original: geometric sampling

Original: -umaze with

Original: R OLLIN on an oracle curriculum in Figure 2.

Original: The confidence interval represents the standard error computed over 8 random seeds.